def trainBatch(net, optimizer): # print('train_iter: ', train_iter) # print('len_train_iter: ', len(train_iter)) # print(type(train_iter.next())) data = train_iter.next() cpu_images, cpu_texts = data batch_size = cpu_images.size(0) utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts) utils.loadData(text, t) utils.loadData(length, l) preds = crnn(image) preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size)) H, cost = seg_ctc_ent_cost(preds, text, preds_size, length, uni_rate=opt.uni_rate) h_cost = (1 - opt.h_rate) * cost - opt.h_rate * H cost_sum = h_cost.data.sum() inf = float("inf") if cost_sum == inf or cost_sum == -inf or cost_sum > 200 * batch_size: print("Warning: received an inf loss, setting loss value to 0") return torch.zeros(H.size()), torch.zeros(cost.size()), torch.zeros( h_cost.size()) crnn.zero_grad() h_cost.backward() torch.nn.utils.clip_grad_norm(crnn.parameters(), opt.max_norm) optimizer.step() return H / batch_size, cost / batch_size, h_cost / batch_size
def val(net, dataset, max_iter=100): print('Start val') for p in crnn.parameters(): p.requires_grad = False net.eval() data_loader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=opt.batchSize, num_workers=int(opt.workers)) val_iter = iter(data_loader) i = 0 n_correct = 0 # loss averager avg_h_val = utils.averager() avg_cost_val = utils.averager() avg_h_cost_val = utils.averager() if opt.eval_all: max_iter = len(data_loader) else: max_iter = min(max_iter, len(data_loader)) for i in range(max_iter): data = val_iter.next() # print('data: ', data) # print(data) i += 1 cpu_images, cpu_texts = data batch_size = cpu_images.size(0) utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts) utils.loadData(text, t) utils.loadData(length, l) # print('len(cpu_images): ', len(cpu_images)) # print('cpu_texts ', cpu_texts) # print('len(cpu_texts): ', len(cpu_texts)) # print('l ', l) # print(len(l)) # print('length ', length) preds = crnn(image) # size = 26, 64, 96 # print(preds.size()) preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size)) H, cost = seg_ctc_ent_cost(preds, text, preds_size, length, uni_rate=opt.uni_rate) h_cost = (1 - opt.h_rate) * cost - opt.h_rate * H avg_h_val.add(H / batch_size) avg_cost_val.add(cost / batch_size) avg_h_cost_val.add(h_cost / batch_size) _, preds = preds.max(2) # size = 26, 64 # print(preds.size()) # preds = preds.squeeze(2) preds = preds.transpose(1, 0).contiguous().view(-1) sim_preds = converter.decode(preds.data, preds_size.data, raw=False) for idx, (pred, target) in enumerate(zip(sim_preds, cpu_texts)): if pred == target.lower(): 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('%-30s => %-30s, gt: %-30s' % (raw_pred, pred, gt)) accuracy = n_correct / float(max_iter * opt.batchSize) print( 'Test H: %f, Cost: %f, H Cost: %f, accuray: %f' % (avg_h_val.val(), avg_cost_val.val(), avg_h_cost_val.val(), accuracy))