def val(net, criterion, max_iter=10):
print('Validating...')
for p in crnn.parameters():
p.requires_grad = False
net.eval()
n_correct = 0
loss_avg = utils.averager()
for i, data in enumerate(test_loader):
cpu_images, cpu_texts = data
batchSize = cpu_images.size(0)
txts, lengths = converter.encode(cpu_texts)
# image = Variable( cpu_images )
# text = Variable( txts )
# length = Variable( lengths )
loadData(image, cpu_images)
loadData(length, lengths)
loadData(text, txts)
preds = crnn(image)
preds_size = Variable(torch.IntTensor([preds.size(0)] * preds.size(1)))
cost = criterion(preds, text, preds_size, length)
loss_avg.add(cost)
_, preds = preds.max(2)
# preds = preds.squeeze(2)
preds = preds.transpose(1, 0).contiguous().view(-1)
sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
for pred, target in zip(sim_preds, cpu_texts):
if pred == target:
n_correct += 1
raw_preds = converter.decode(preds.data, preds_size.data,
raw=True)[:opt.n_test_disp]
for raw_pred, pred, gt in zip(raw_preds, sim_preds, cpu_texts):
print('%-20s => %-20s, gt: %-20s' % (raw_pred, pred, gt))
accuracy = n_correct / float(len(test_loader) * opt.batchSize)
print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy))
def main(epoch):
for i, batchData in enumerate(train_loader):
for p in crnn.parameters():
p.requires_grad = True
crnn.train()
cost = trainBatch(batchData)
loss_avg.add(cost)
i += 1
if i % opt.displayInterval == 0:
print('[ %s ][%d/%d][%d/%d] Loss: %f' %
(datetime.now().strftime('%H-%M-%S'), epoch, opt.niter, i,
len(train_loader), loss_avg.val()))
loss_avg.reset()
if i % opt.valInterval == 0:
val(crnn, criterion)
if i % opt.saveInterval == 0:
saveModel(epoch)
cudnn.benchmark = True
crnn.cuda()
crnn = torch.nn.DataParallel(crnn, device_ids=range(opt.ngpu))
image = image.cuda()
criterion = criterion.cuda()
image = Variable(image)
text = Variable(text)
length = Variable(length)
# loss averager
loss_avg = utils.averager()
# setup optimizer
if opt.adam:
optimizer = optim.Adam(crnn.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
elif opt.adadelta:
optimizer = optim.Adadelta(crnn.parameters(), lr=opt.lr)
else:
optimizer = optim.RMSprop(crnn.parameters(), lr=opt.lr)
def loadData(v, data):
v.data.resize_(data.size()).copy_(data)
def val(net, criterion, max_iter=10):
print('Validating...')