def __init__(self):
encoder_path = './expr/attentioncnn/encoder_600.pth'
decoder_path = './expr/attentioncnn/decoder_600.pth'
self.alphabet = '0123456789'
self.max_length = 7 # 最长字符串的长度
self.EOS_TOKEN = 1
self.use_gpu = True
self.max_width = 220
self.converter = utils.strLabelConverterForAttention(self.alphabet)
self.transform = transforms.ToTensor()
nclass = len(self.alphabet) + 3
encoder = crnn.CNN(32, 1, 256) # 编码器
decoder = crnn.decoder(256,
nclass) # seq to seq的解码器, nclass在decoder中还加了2
if encoder_path and decoder_path:
print('loading pretrained models ......')
encoder.load_state_dict(torch.load(encoder_path))
decoder.load_state_dict(torch.load(decoder_path))
if torch.cuda.is_available() and self.use_gpu:
encoder = encoder.cuda()
decoder = decoder.cuda()
self.encoder = encoder.eval()
self.decoder = decoder.eval()
shuffle=False,
sampler=sampler,
num_workers=int(opt.workers),
collate_fn=dataset.alignCollate(
imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio=opt.keep_ratio))
test_dataset = dataset.listDataset(list_file=opt.vallist,
transform=dataset.resizeNormalize(
(opt.imgW, opt.imgH)))
nclass = len(alphabet) + 3 # decoder的时候，需要的类别数,3 for SOS,EOS和blank
nc = 1
converter = utils.strLabelConverterForAttention(alphabet)
# criterion = torch.nn.CrossEntropyLoss()
criterion = torch.nn.NLLLoss()
def weights_init(model):
# Official init from torch repo.
for m in model.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.constant_(m.bias, 0)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=opt.batchSize,
shuffle=False, sampler=sampler,
num_workers=int(opt.workers),
collate_fn=dataset.alignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio=opt.keep_ratio))
print (len(train_loader))
#test_dataset = dataset.listDataset(list_file =opt.vallist, transform=dataset.resizeNormalize((100, 32)))
test_dataset = dataset.listDataset(list_file =opt.vallist)
alphabet = str1.decode('utf-8')
nclass = len(alphabet)
nc = 1
converterAttention = utils.strLabelConverterForAttention(alphabet, opt.sep)
converterCTC = utils.strLabelConverterForCTC(alphabet, opt.sep)
criterionAttention = torch.nn.CrossEntropyLoss()
criterionCTC = CTCLoss()
# custom weights initialization called on crnn
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)