mirror of
https://github.com/netfun2000/dddd_trainer.git
synced 2026-10-14 20:53:25 +08:00
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
'''
|
|
不记得从哪套模型改的了，可能来自于部门mobildenetv2
|
|
'''
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class DdddOcr(nn.Module):
|
|
def __init__(self, nc=3, leakyRelu=False):
|
|
super(DdddOcr, self).__init__()
|
|
# assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
|
|
|
|
ks = [3, 3, 3, 3, 3, 3, 2]
|
|
ps = [1, 1, 1, 1, 1, 1, 0]
|
|
ss = [1, 1, 1, 1, 1, 1, 1]
|
|
nm = [16, 32, 64, 64, 128, 128, 128]
|
|
|
|
cnn = nn.Sequential()
|
|
|
|
def convRelu(i, batchNormalization=False):
|
|
nIn = nc if i == 0 else nm[i - 1]
|
|
nOut = nm[i]
|
|
cnn.add_module('conv{0}'.format(i),
|
|
nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i]))
|
|
if batchNormalization:
|
|
cnn.add_module('batchnorm{0}'.format(i), nn.BatchNorm2d(nOut))
|
|
if leakyRelu:
|
|
cnn.add_module('relu{0}'.format(i),
|
|
nn.LeakyReLU(0.2, inplace=True))
|
|
else:
|
|
cnn.add_module('relu{0}'.format(i), nn.ReLU(True))
|
|
|
|
convRelu(0)
|
|
cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2)) # 64x16x64
|
|
convRelu(1)
|
|
cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2)) # 128x8x32
|
|
convRelu(2, True)
|
|
convRelu(3)
|
|
cnn.add_module('pooling{0}'.format(2),
|
|
nn.MaxPool2d((2, 2), (2, 1), (0, 1))) # 256x4x16
|
|
convRelu(4, True)
|
|
convRelu(5)
|
|
cnn.add_module('pooling{0}'.format(3),
|
|
nn.MaxPool2d((2, 2), (2, 1), (0, 1))) # 512x2x16
|
|
convRelu(6, True) # 512x1x16
|
|
|
|
self.cnn = cnn
|
|
|
|
def forward(self, input):
|
|
return self.cnn(input)
|
|
|
|
def test():
|
|
net = DdddOcr(1)
|
|
x = torch.randn(1, 1, 64, 224)
|
|
y = net(x)
|
|
print(y.size())
|
|
y = y.permute(3, 0, 1, 2)
|
|
w, b, c, h = y.shape
|
|
y = y.view(w, b, c * h)
|
|
print(y.size())
|
|
|
|
if __name__ == '__main__':
|
|
test() |