def train(opt): """ dataset preparation """ opt.select_data = opt.select_data.split('-') opt.batch_ratio = opt.batch_ratio.split('-') train_dataset = Batch_Balanced_Dataset(opt) AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD) valid_dataset = hierarchical_dataset(root=opt.valid_data, opt=opt) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=opt.batch_size, shuffle=True, # 'True' to check training progress with validation function. num_workers=int(opt.workers), collate_fn=AlignCollate_valid, pin_memory=True) print('-' * 80) """ model configuration """ if 'CTC' in opt.Prediction: converter = CTCLabelConverter(opt.character) else: converter = AttnLabelConverter(opt.character) opt.num_class = len(converter.character) if opt.rgb: opt.input_channel = 3 model = Model(opt) print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel, opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction, opt.SequenceModeling, opt.Prediction) # weight initialization for name, param in model.named_parameters(): if 'localization_fc2' in name: print(f'Skip {name} as it is already initialized') continue try: if 'bias' in name: init.constant_(param, 0.0) elif 'weight' in name: init.kaiming_normal_(param) except Exception as e: # for batchnorm. if 'weight' in name: param.data.fill_(1) continue # data parallel for multi-GPU model = torch.nn.DataParallel(model).cuda() model.train() if opt.continue_model != '': if opt.without_prediction: load_model_without_prediction(opt.continue_model, model) print(f'loading pretrained model from {opt.continue_model}, without prediction layer') else: print(f'loading pretrained model from {opt.continue_model}') model.load_state_dict(torch.load(opt.continue_model)) print("Model:") print(model) """ setup loss """ if 'CTC' in opt.Prediction: criterion = torch.nn.CTCLoss(zero_infinity=True).cuda() else: criterion = torch.nn.CrossEntropyLoss(ignore_index=0).cuda() # ignore [GO] token = ignore index 0 # loss averager loss_avg = Averager() # filter that only require gradient decent filtered_parameters = [] params_num = [] for p in filter(lambda p: p.requires_grad, model.parameters()): filtered_parameters.append(p) params_num.append(np.prod(p.size())) print('Trainable params num : ', sum(params_num)) # [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())] # setup optimizer if opt.adam: optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999)) else: optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps) print("Optimizer:") print(optimizer) """ final options """ # print(opt) with open(f'./saved_models/{opt.experiment_name}/opt.txt', 'a') as opt_file: opt_log = '------------ Options -------------\n' args = vars(opt) for k, v in args.items(): opt_log += f'{str(k)}: {str(v)}\n' opt_log += '---------------------------------------\n' print(opt_log) opt_file.write(opt_log) """ start training """ start_iter = 0 if opt.continue_model != '': print(f'continue to train, start_iter: {start_iter}') start_time = time.time() best_accuracy = -1 best_norm_ED = 1e+6 i = start_iter while True: # train part for p in model.parameters(): p.requires_grad = True image_tensors, labels = train_dataset.get_batch() image = image_tensors.cuda() text, length = converter.encode(labels) batch_size = image.size(0) if 'CTC' in opt.Prediction: preds = model(image, text).log_softmax(2) preds_size = torch.IntTensor([preds.size(1)] * batch_size) preds = preds.permute(1, 0, 2) # to use CTCLoss format cost = criterion(preds, text, preds_size, length) else: preds = model(image, text) target = text[:, 1:] # without [GO] Symbol cost = criterion(preds.view(-1, preds.shape[-1]), target.contiguous().view(-1)) model.zero_grad() cost.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip) # gradient clipping with 5 (Default) optimizer.step() loss_avg.add(cost) # validation part if i % opt.valInterval == 0: elapsed_time = time.time() - start_time logging.info(f'[{i}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}') # for log with open(f'./saved_models/{opt.experiment_name}/log_train.txt', 'a') as log: log.write(f'[{i}/{opt.num_iter}] Loss: {loss_avg.val():0.5f} elapsed_time: {elapsed_time:0.5f}\n') loss_avg.reset() model.eval() valid_loss, current_accuracy, current_norm_ED, preds, labels, infer_time, length_of_data = validation( model, criterion, valid_loader, converter, opt) model.train() for pred, gt in zip(preds[:5], labels[:5]): if 'Attn' in opt.Prediction: pred = pred[:pred.find('[s]')] gt = gt[:gt.find('[s]')] print(f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}') log.write(f'{pred:20s}, gt: {gt:20s}, {str(pred == gt)}\n') valid_log = f'[{i}/{opt.num_iter}] valid loss: {valid_loss:0.5f}' valid_log += f' accuracy: {current_accuracy:0.3f}, norm_ED: {current_norm_ED:0.2f}' log.write(valid_log + '\n') # keep best accuracy model if current_accuracy > best_accuracy: best_accuracy = current_accuracy torch.save(model.state_dict(), f'./saved_models/{opt.experiment_name}/mtl_best_accuracy.pth') if current_norm_ED < best_norm_ED: best_norm_ED = current_norm_ED torch.save(model.state_dict(), f'./saved_models/{opt.experiment_name}/best_norm_ED.pth') best_model_log = f'best_accuracy: {best_accuracy:0.3f}, best_norm_ED: {best_norm_ED:0.2f}' logging.info(best_model_log) log.write(best_model_log + '\n') # save model per 1e+5 iter. if (i + 1) % 50000 == 0: torch.save( model.state_dict(), f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth') if i == opt.num_iter: logging.info('end the training') sys.exit() i += 1
def train(opt): #logging.info(opt) train_dataset = Batch_Dataset(opt) AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD) valid_dataset = LmdbDataset(root=opt.valid_data, opt=opt) valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=int(opt.workers), collate_fn=AlignCollate_valid, pin_memory=True) print('-' * 80) """ model configuration """ ctc_converter = CTCLabelConverter(opt.character, opt.subword) attn_converter = AttnLabelConverter(opt.character, opt.subword, opt.batch_max_length) opt.num_class = len(attn_converter.character) opt.ctc_num_class = len(ctc_converter.character) print("ctc num class {}".format(len(ctc_converter.character))) print("attention num class {}".format(len(attn_converter.character))) if opt.rgb: opt.input_channel = 3 model = MyModel(opt) # weight initialization for name, param in model.named_parameters(): if 'localization_fc2' in name: print('Skip {name} as it is already initialized'.format(name)) continue try: if 'bias' in name: init.constant_(param, 0.0) elif 'weight' in name: init.kaiming_normal_(param) except Exception as e: if 'weight' in name: param.data.fill_(1) continue model = torch.nn.DataParallel(model).to(device) model.train() if opt.continue_model != '': print('loading pretrained model from {}'.format(opt.continue_model)) model.load_state_dict(torch.load(opt.continue_model)) """ setup loss """ ctc_criterion = torch.nn.CTCLoss(zero_infinity=True).to(device) attn_criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device) loss_avg = Averager() filtered_parameters = [] params_num = [] for p in filter(lambda p: p.requires_grad, model.parameters()): filtered_parameters.append(p) params_num.append(np.prod(p.size())) print('Trainable params num : ', sum(params_num)) if opt.adam: optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999)) else: optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps) print("Optimizer:") print(optimizer) """ final options """ with open(osj(opt.outPath, '{}/opt.txt'.format(opt.experiment_name)), 'a') as opt_file: opt_log = '------------ Options -------------\n' args = vars(opt) for k, v in args.items(): opt_log += '{}: {}\n'.format(str(k), str(v)) opt_log += '---------------------------------------\n' print(opt_log) opt_file.write(opt_log) """ start training """ start_iter = 0 if opt.continue_model != '': print('continue to train, start_iter: {}'.format(start_iter)) start_time = time.time() best_accuracy = -1 i = start_iter while True: # train part for p in model.parameters(): p.requires_grad = True image_tensors, labels = train_dataset.get_batch() image = image_tensors.to(device) ctc_text, ctc_length = ctc_converter.encode(labels) attn_text, attn_length = attn_converter.encode(labels) batch_size = image.size(0) # ctc loss ctc_preds, attn_preds = model(image, attn_text) ctc_preds = ctc_preds.log_softmax(2) preds_size = torch.IntTensor([ctc_preds.size(1)] * batch_size) ctc_preds = ctc_preds.permute(1, 0, 2) ctc_cost = ctc_criterion(ctc_preds, ctc_text, preds_size, ctc_length) # attn loss target = attn_text[:, 1:] attn_cost = attn_criterion(attn_preds.view(-1, attn_preds.shape[-1]), target.contiguous().view(-1)) cost = opt.ctc_weight * ctc_cost + (1.0 - opt.ctc_weight) * attn_cost model.zero_grad() cost.backward() torch.nn.utils.clip_grad_norm_( model.parameters(), opt.grad_clip) # gradient clipping with 5 (Default) optimizer.step() loss_avg.add(cost) # validation part if i % opt.valInterval == 0: elapsed_time = time.time() - start_time logging.info('[{}/{}] Loss: {:0.5f} elapsed_time: {:0.5f}'.format( i, opt.num_iter, loss_avg.val(), elapsed_time)) # for log with open( osj(opt.outPath, '{}/log_train.txt'.format(opt.experiment_name)), 'a') as log: log.write( '[{}/{}] Loss: {:0.5f} elapsed_time: {:0.5f}\n'.format( i, opt.num_iter, loss_avg.val(), elapsed_time)) loss_avg.reset() model.eval() with torch.no_grad(): valid_loss, current_accuracy, ctc_accuracy, current_norm_ED, preds, labels, infer_time, length_of_data \ = mtl_validation(model, ctc_criterion, attn_criterion, valid_loader, ctc_converter, attn_converter, opt) model.train() for pred, gt in zip(preds[:5], labels[:5]): pred = pred[:pred.find('[s]')] gt = gt[:gt.find('[s]')] print('{:20s}, gt: {:20s}, {}'.format( pred, gt, str(pred == gt))) log.write('{:20s}, gt: {:20s}, {}\n'.format( pred, gt, str(pred == gt))) valid_log = '[{}/{}] valid loss: {:0.5f}'.format( i, opt.num_iter, valid_loss) valid_log += ' accuracy: {:0.3f}'.format(current_accuracy) log.write(valid_log + '\n') # save best accuracy model if current_accuracy > best_accuracy: best_accuracy = current_accuracy torch.save( model.state_dict(), osj(opt.outPath, '{}/best_accuracy.pth'.format( opt.experiment_name))) best_model_log = 'best_accuracy: {:0.3f}'.format(best_accuracy) logging.info(best_model_log) log.write(best_model_log + '\n') if (i + 1) % 50000 == 0: torch.save( model.state_dict(), osj(opt.outPath, '{}/iter_{}.pth'.format(opt.experiment_name, i + 1))) if i == opt.num_iter: logging.info('end the training') sys.exit() i += 1