def train(loader, model, crit, optimizer, lr_scheduler, opt, rl_crit=None): model.train() model = nn.DataParallel(model) for epoch in range(opt["epochs"]): lr_scheduler.step() iteration = 0 # If start self crit training if opt["self_crit_after"] != -1 and epoch >= opt["self_crit_after"]: sc_flag = True init_cider_scorer(opt["cached_tokens"]) else: sc_flag = False for data in loader: torch.cuda.synchronize() fc_feats = Variable(data['fc_feats']).cuda() labels = Variable(data['labels']).long().cuda() masks = Variable(data['masks']).cuda() optimizer.zero_grad() if not sc_flag: seq_probs, _ = model(fc_feats, labels, 'train') loss = crit(seq_probs, labels[:, 1:], masks[:, 1:]) else: seq_probs, seq_preds = model(fc_feats, mode='inference', opt=opt) reward = get_self_critical_reward(model, fc_feats, data, seq_preds) print(reward.shape) loss = rl_crit( seq_probs, seq_preds, Variable(torch.from_numpy(reward).float().cuda())) loss.backward() utils.clip_gradient(optimizer, opt["grad_clip"]) optimizer.step() train_loss = loss.data[0] torch.cuda.synchronize() iteration += 1 if not sc_flag: print("iter %d (epoch %d), train_loss = %.6f" % (iteration, epoch, train_loss)) else: print("iter %d (epoch %d), avg_reward = %.6f" % (iteration, epoch, np.mean(reward[:, 0]))) if epoch != 0 and epoch % opt["save_checkpoint_every"] == 0: model_path = os.path.join(opt["checkpoint_path"], 'model_%d.pth' % (epoch)) model_info_path = os.path.join(opt["checkpoint_path"], 'model_score.txt') torch.save(model.state_dict(), model_path) print("model saved to %s" % (model_path)) with open(model_info_path, 'a') as f: f.write("model_%d, loss: %.6f\n" % (epoch, train_loss))
def train(opt): opt.use_att = utils.if_use_att(opt.caption_model) loader = DataLoader(opt) opt.vocab_size = loader.vocab_size opt.seq_length = loader.seq_length tf_summary_writer = tf and tf.summary.FileWriter(opt.checkpoint_path) infos = {} histories = {} if opt.start_from is not None: # open old infos and check if models are compatible with open(os.path.join(opt.start_from, 'infos_' + opt.id + '.pkl')) as f: infos = cPickle.load(f) saved_model_opt = infos['opt'] need_be_same = [ "caption_model", "rnn_type", "rnn_size", "num_layers" ] for checkme in need_be_same: assert vars(saved_model_opt)[checkme] == vars( opt )[checkme], "Command line argument and saved model disagree on '%s' " % checkme if os.path.isfile( os.path.join(opt.start_from, 'histories_' + opt.id + '.pkl')): with open( os.path.join(opt.start_from, 'histories_' + opt.id + '.pkl')) as f: histories = cPickle.load(f) iteration = infos.get('iter', 0) epoch = infos.get('epoch', 0) val_result_history = histories.get('val_result_history', {}) loss_history = histories.get('loss_history', {}) lr_history = histories.get('lr_history', {}) ss_prob_history = histories.get('ss_prob_history', {}) loader.iterators = infos.get('iterators', loader.iterators) loader.split_ix = infos.get('split_ix', loader.split_ix) if opt.load_best_score == 1: best_val_score = infos.get('best_val_score', None) model = models.setup(opt) model.cuda() update_lr_flag = True # Assure in training mode model.train() crit = utils.LanguageModelCriterion() rl_crit = utils.RewardCriterion() optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay) # Load the optimizer if vars(opt).get('start_from', None) is not None and os.path.isfile( os.path.join(opt.start_from, "optimizer.pth")): optimizer.load_state_dict( torch.load(os.path.join(opt.start_from, 'optimizer.pth'))) while True: if update_lr_flag: # Assign the learning rate if epoch > opt.learning_rate_decay_start and opt.learning_rate_decay_start >= 0: frac = (epoch - opt.learning_rate_decay_start ) // opt.learning_rate_decay_every decay_factor = opt.learning_rate_decay_rate**frac opt.current_lr = opt.learning_rate * decay_factor else: opt.current_lr = opt.learning_rate utils.set_lr(optimizer, opt.current_lr) # Assign the scheduled sampling prob if epoch > opt.scheduled_sampling_start and opt.scheduled_sampling_start >= 0: frac = (epoch - opt.scheduled_sampling_start ) // opt.scheduled_sampling_increase_every opt.ss_prob = min(opt.scheduled_sampling_increase_prob * frac, opt.scheduled_sampling_max_prob) model.ss_prob = opt.ss_prob # If start self critical training if opt.self_critical_after != -1 and epoch >= opt.self_critical_after: sc_flag = True init_cider_scorer(opt.cached_tokens) else: sc_flag = False update_lr_flag = False start = time.time() # Load data from train split (0) data = loader.get_batch('train') print('Read data:', time.time() - start) torch.cuda.synchronize() start = time.time() tmp = [ data['fc_feats'], data['att_feats'], data['labels'], data['masks'] ] tmp = [ Variable(torch.from_numpy(_), requires_grad=False).cuda() for _ in tmp ] fc_feats, att_feats, labels, masks = tmp optimizer.zero_grad() if not sc_flag: loss = crit(model(fc_feats, att_feats, labels), labels[:, 1:], masks[:, 1:]) else: gen_result, sample_logprobs = model.sample(fc_feats, att_feats, {'sample_max': 0}) reward = get_self_critical_reward(model, fc_feats, att_feats, data, gen_result) loss = rl_crit( sample_logprobs, gen_result, Variable(torch.from_numpy(reward).float().cuda(), requires_grad=False)) loss.backward() utils.clip_gradient(optimizer, opt.grad_clip) optimizer.step() train_loss = loss.data[0] torch.cuda.synchronize() end = time.time() if not sc_flag: print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \ .format(iteration, epoch, train_loss, end - start)) else: print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \ .format(iteration, epoch, np.mean(reward[:,0]), end - start)) # Update the iteration and epoch iteration += 1 if data['bounds']['wrapped']: epoch += 1 update_lr_flag = True # Write the training loss summary if (iteration % opt.losses_log_every == 0): if tf is not None: add_summary_value(tf_summary_writer, 'train_loss', train_loss, iteration) add_summary_value(tf_summary_writer, 'learning_rate', opt.current_lr, iteration) add_summary_value(tf_summary_writer, 'scheduled_sampling_prob', model.ss_prob, iteration) if sc_flag: add_summary_value(tf_summary_writer, 'avg_reward', np.mean(reward[:, 0]), iteration) tf_summary_writer.flush() loss_history[iteration] = train_loss if not sc_flag else np.mean( reward[:, 0]) lr_history[iteration] = opt.current_lr ss_prob_history[iteration] = model.ss_prob # make evaluation on validation set, and save model if (iteration % opt.save_checkpoint_every == 0): # eval model eval_kwargs = {'split': 'val', 'dataset': opt.input_json} eval_kwargs.update(vars(opt)) val_loss, predictions, lang_stats = eval_utils.eval_split( model, crit, loader, eval_kwargs) # Write validation result into summary if tf is not None: add_summary_value(tf_summary_writer, 'validation loss', val_loss, iteration) if lang_stats is not None: for k, v in lang_stats.items(): add_summary_value(tf_summary_writer, k, v, iteration) tf_summary_writer.flush() val_result_history[iteration] = { 'loss': val_loss, 'lang_stats': lang_stats, 'predictions': predictions } # Save model if is improving on validation result if opt.language_eval == 1: current_score = lang_stats['CIDEr'] else: current_score = -val_loss best_flag = False if True: # if true if best_val_score is None or current_score > best_val_score: best_val_score = current_score best_flag = True checkpoint_path = os.path.join(opt.checkpoint_path, 'model.pth') torch.save(model.state_dict(), checkpoint_path) print("model saved to {}".format(checkpoint_path)) optimizer_path = os.path.join(opt.checkpoint_path, 'optimizer.pth') torch.save(optimizer.state_dict(), optimizer_path) # Dump miscalleous informations infos['iter'] = iteration infos['epoch'] = epoch infos['iterators'] = loader.iterators infos['split_ix'] = loader.split_ix infos['best_val_score'] = best_val_score infos['opt'] = opt infos['vocab'] = loader.get_vocab() histories['val_result_history'] = val_result_history histories['loss_history'] = loss_history histories['lr_history'] = lr_history histories['ss_prob_history'] = ss_prob_history with open( os.path.join(opt.checkpoint_path, 'infos_' + opt.id + '.pkl'), 'wb') as f: cPickle.dump(infos, f) with open( os.path.join(opt.checkpoint_path, 'histories_' + opt.id + '.pkl'), 'wb') as f: cPickle.dump(histories, f) if best_flag: checkpoint_path = os.path.join(opt.checkpoint_path, 'model-best.pth') torch.save(model.state_dict(), checkpoint_path) print("model saved to {}".format(checkpoint_path)) with open( os.path.join(opt.checkpoint_path, 'infos_' + opt.id + '-best.pkl'), 'wb') as f: cPickle.dump(infos, f) # Stop if reaching max epochs if epoch >= opt.max_epochs and opt.max_epochs != -1: break
def train(dataset, loader, model, rem, crit, optimizer, lr_scheduler, opt, rl_crit=None): writer = SummaryWriter('./runs/video_caption22') model.load_state_dict( torch.load( '/home/diml/video-caption.pytorch/save/RECON222_model_200.pth')) rem.load_state_dict( torch.load( '/home/diml/video-caption.pytorch/save/RECON222_module_200.pth')) #model.load_state_dict(torch.load('/home/diml/video-caption.pytorch/save/new_model_200.pth')) #model = nn.DataParallel(model) model.train() rem.train() vocab = dataset.get_vocab() for epoch in trange(opt["epochs"]): t_loss = [0, 0, 0] # ============================================================================= # model.eval() # ev.demov(model,crit, dataset, dataset.get_vocab(),opt) # ============================================================================= lr_scheduler.step() iteration = 0 # If start self crit training if opt["self_crit_after"] != -1 and epoch >= opt["self_crit_after"]: sc_flag = True init_cider_scorer(opt["cached_tokens"]) else: sc_flag = False for idx, data in enumerate(loader): torch.cuda.synchronize() fc_feats = data['fc_feats'].cuda() labels = data['labels'].cuda() labels2 = data['labels2'].cuda() masks2 = data['masks2'].cuda() masks = data['masks'].cuda() optimizer.zero_grad() if not sc_flag: seq_probs, seq_preds, hn, de_hn = model( fc_feats, labels, 'train') loss_C = crit(seq_probs, labels[:, 1:], masks[:, 1:]) fake_en_hn = rem(de_hn, seq_probs) f_seq_probs, f_seq_preds, hn, de_hn = model(fc_feats, labels2, 'train', h=fake_en_hn) loss_R = crit(f_seq_probs, labels2[:, 1:], masks2[:, 1:]) loss = loss_R + loss_C else: seq_probs, seq_preds = model(fc_feats, mode='inference', opt=opt) reward = get_self_critical_reward(model, fc_feats, data, seq_preds) print(reward.shape) loss = rl_crit(seq_probs, seq_preds, torch.from_numpy(reward).float().cuda()) t_loss[0] += loss.item() t_loss[1] += loss_C.item() t_loss[2] += loss_R.item() loss.backward() clip_grad_value_(model.parameters(), opt['grad_clip']) optimizer.step() train_loss = loss.item() torch.cuda.synchronize() iteration += 1 if not sc_flag: print("iter %d (epoch %d), train_loss = %.6f" % (iteration, epoch, train_loss)) else: print("iter %d (epoch %d), avg_reward = %.6f" % (iteration, epoch, np.mean(reward[:, 0]))) writer.add_scalar('training total loss', t_loss[0] / 140, epoch + 200) writer.add_scalar('training Caption loss', t_loss[1] / 140, epoch + 200) writer.add_scalar('training Reconstruction loss', t_loss[2] / 140, epoch + 200) if epoch % opt["save_checkpoint_every"] == 0: model_path = os.path.join(opt["checkpoint_path"], 'RECON222_model_%d.pth' % (epoch + 200)) rem_path = os.path.join(opt["checkpoint_path"], 'RECON222_module_%d.pth' % (epoch + 200)) model_info_path = os.path.join(opt["checkpoint_path"], 'RECON222_model_score.txt') torch.save(model.state_dict(), model_path) torch.save(rem.state_dict(), rem_path) print("model saved to %s" % (model_path)) with open(model_info_path, 'a') as f: f.write("model_%d, loss: %.6f\n" % (epoch, train_loss)) with torch.no_grad(): _, seq_preds, __, ___ = model(fc_feats, mode='inference', opt=opt) _, f_seq_preds, __, ___ = model(fc_feats, mode='inference', h=fake_en_hn, opt=opt) origin = utils.decode_sequence(vocab, seq_preds)[0] revision = utils.decode_sequence(vocab, f_seq_preds)[0] with open('./results/training_versus.txt', 'a') as f: f.write("epoch is %d \n" % epoch) origin = "origin caption: " + origin + "\n" revision = "revision caption: " + revision + "\n" f.write(origin) f.write(revision)
def train(dataset, loader, model, crit, optimizer, lr_scheduler, opt, rl_crit=None): writer = SummaryWriter('./runs/video_caption_basic') model.load_state_dict( torch.load('/home/diml/video-caption.pytorch/save/new_model_200.pth')) #model = nn.DataParallel(model) model.train() vocab = dataset.get_vocab() for epoch in trange(300): t_loss = 0 # ============================================================================= # model.eval() # ev.demov(model,crit, dataset, dataset.get_vocab(),opt) # ============================================================================= lr_scheduler.step() iteration = 0 # If start self crit training if opt["self_crit_after"] != -1 and epoch >= opt["self_crit_after"]: sc_flag = True init_cider_scorer(opt["cached_tokens"]) else: sc_flag = False for idx, data in enumerate(loader): torch.cuda.synchronize() fc_feats = data['fc_feats'].cuda() labels = data['labels'].cuda() masks = data['masks'].cuda() optimizer.zero_grad() if not sc_flag: seq_probs, seq_preds, hn, de_hn = model( fc_feats, labels, 'train') loss_C = crit(seq_probs, labels[:, 1:], masks[:, 1:]) loss = loss_C else: seq_probs, seq_preds = model(fc_feats, mode='inference', opt=opt) reward = get_self_critical_reward(model, fc_feats, data, seq_preds) print(reward.shape) loss = rl_crit(seq_probs, seq_preds, torch.from_numpy(reward).float().cuda()) t_loss += loss.item() loss.backward() clip_grad_value_(model.parameters(), opt['grad_clip']) optimizer.step() train_loss = loss.item() torch.cuda.synchronize() iteration += 1 if not sc_flag: print("iter %d (epoch %d), train_loss = %.6f" % (iteration, epoch, train_loss)) else: print("iter %d (epoch %d), avg_reward = %.6f" % (iteration, epoch + 201, np.mean(reward[:, 0]))) writer.add_scalar('training total loss', t_loss / 140, epoch + 200) if epoch % opt["save_checkpoint_every"] == 0: model_path = os.path.join(opt["checkpoint_path"], 'new_model_%d.pth' % (epoch + 200)) model_info_path = os.path.join(opt["checkpoint_path"], 'Rnew_model_score.txt') torch.save(model.state_dict(), model_path) print("model saved to %s" % (model_path)) with open(model_info_path, 'a') as f: f.write("model_%d, loss: %.6f\n" % (epoch, train_loss)) with torch.no_grad(): _, seq_preds, __, ___ = model(fc_feats, mode='inference', opt=opt) print(utils.decode_sequence(vocab, seq_preds)[0])
def train(loader, model, crit, optimizer, lr_scheduler, opt, rl_crit=None): model.train() if opt['visdom']: viz = visdom.Visdom(env='train') loss_win = viz.line(np.arange(1), opts={'title': 'loss'}) for epoch in range(opt["epochs"]): lr_scheduler.step() iteration = 0 # If start self crit training # print(opt["self_crit_after"]) if opt["self_crit_after"] != -1 and epoch >= opt[ "self_crit_after"]: #每多少次保存一下 sc_flag = True init_cider_scorer(opt["cached_tokens"]) else: sc_flag = False # print(model) for data in loader: # print(data) torch.cuda.synchronize() fc_feats = data['fc_feats'].cuda() # voice_feats = data['voice_feats'].cuda() if opt['with_hand'] == 1: hand_feats = data['hand_feats'].cuda() hand_pro = data['hand_pro'].cuda() labels = data['labels'].cuda() masks = data['masks'].cuda() #print(sc_flag) optimizer.zero_grad() if not sc_flag: # seq_probs, _ = model(fc_feats, voice_feats, hand_feats, labels, 'train') if opt['with_hand'] == 1: seq_probs, _ = model(fc_feats, hand_feats, hand_pro, labels, 'train') else: seq_probs, _ = model.forward2(fc_feats, labels, 'train') loss = crit(seq_probs, labels[:, 1:], masks[:, 1:]) # todo 下面else部分没有修改声音和手语的内容 else: seq_probs, seq_preds = model(fc_feats, mode='inference', opt=opt) reward = get_self_critical_reward(model, fc_feats, data, seq_preds) print(reward.shape) loss = rl_crit(seq_probs, seq_preds, torch.from_numpy(reward).float().cuda()) loss.backward() clip_grad_value_(model.parameters(), opt['grad_clip']) optimizer.step() train_loss = loss.item() torch.cuda.synchronize() iteration += 1 if not sc_flag: print("?iter %d (epoch %d), train_loss = %.6f" % (iteration, epoch, train_loss)) if opt['visdom']: viz.line(Y=np.array([train_loss]), X=np.array([epoch]), win=loss_win, update='append') else: print("??iter %d (epoch %d), avg_reward = %.6f" % (iteration, epoch, np.mean(reward[:, 0]))) if epoch % opt["save_checkpoint_every"] == 0: model_path = os.path.join(opt["checkpoint_path"], 'model_%d.pth' % (epoch)) model_info_path = os.path.join(opt["checkpoint_path"], 'model_score.txt') torch.save(model.state_dict(), model_path) # print("model saved to %s" % (model_path)) with open(model_info_path, 'a') as f: f.write("model_%d, loss: %.6f\n" % (epoch, train_loss))
def train(train_loader, val_loader, model, crit, optimizer, lr_scheduler, opt, rl_crit=None): model.train() model = nn.DataParallel(model) # lowest val loss best_loss = None for epoch in range(opt.epochs): lr_scheduler.step() iteration = 0 # If start self crit training if opt.self_crit_after != -1 and epoch >= opt.self_crit_after: sc_flag = True init_cider_scorer(opt.cached_tokens) else: sc_flag = False for data in train_loader: torch.cuda.synchronize() fc_feats = Variable(data['fc_feats']).cuda() labels = Variable(data['labels']).long().cuda() masks = Variable(data['masks']).cuda() if not sc_flag: seq_probs, predicts = model(fc_feats, labels) loss = crit(seq_probs, labels[:, 1:], masks[:, 1:]) else: gen_result, sample_logprobs = model.sample(fc_feats, vars(opt)) # print(gen_result) reward = get_self_critical_reward(model, fc_feats, data, gen_result) loss = rl_crit( sample_logprobs, gen_result, Variable(torch.from_numpy(reward).float().cuda())) optimizer.zero_grad() loss.backward() utils.clip_gradient(optimizer, opt.grad_clip) optimizer.step() train_loss = loss.data[0] torch.cuda.synchronize() iteration += 1 if not sc_flag: print("iter %d (epoch %d), train_loss = %.6f" % (iteration, epoch, train_loss)) else: print("iter %d (epoch %d), avg_reward = %.3f" % (iteration, epoch, np.mean(reward[:, 0]))) # lowest val loss if epoch % opt.save_checkpoint_every == 0: checkpoint_path = os.path.join(opt.checkpoint_path, 'model_%d.pth' % (epoch)) torch.save(model.state_dict(), checkpoint_path) print("model saved to %s" % (checkpoint_path)) val_loss = val(val_loader, model, crit) print("Val loss is: %.6f" % (val_loss)) model.train() if best_loss is None or val_loss < best_loss: print("(epoch %d), now lowest val loss is %.6f" % (epoch, val_loss)) checkpoint_path = os.path.join(opt.checkpoint_path, 'model_best.pth') torch.save(model.state_dict(), checkpoint_path) print("best model saved to %s" % (checkpoint_path)) best_loss = val_loss
def train(opt): opt.use_att = utils.if_use_att(opt.caption_model) loader = DataLoader(opt) opt.vocab_size = loader.vocab_size opt.seq_length = loader.seq_length tf_summary_writer = tf and tf.summary.FileWriter(opt.checkpoint_path) infos = {} histories = {} if opt.start_from is not None: # open old infos and check if models are compatible with open(os.path.join(opt.start_from, 'infos_'+opt.id+'.pkl')) as f: infos = cPickle.load(f) saved_model_opt = infos['opt'] need_be_same=["caption_model", "rnn_type", "rnn_size", "num_layers"] for checkme in need_be_same: assert vars(saved_model_opt)[checkme] == vars(opt)[checkme], "Command line argument and saved model disagree on '%s' " % checkme if os.path.isfile(os.path.join(opt.start_from, 'histories_'+opt.id+'.pkl')): with open(os.path.join(opt.start_from, 'histories_'+opt.id+'.pkl')) as f: histories = cPickle.load(f) iteration = infos.get('iter', 0) epoch = infos.get('epoch', 0) val_result_history = histories.get('val_result_history', {}) loss_history = histories.get('loss_history', {}) lr_history = histories.get('lr_history', {}) ss_prob_history = histories.get('ss_prob_history', {}) loader.iterators = infos.get('iterators', loader.iterators) loader.split_ix = infos.get('split_ix', loader.split_ix) if opt.load_best_score == 1: best_val_score = infos.get('best_val_score', None) model = models.setup(opt) model.cuda() update_lr_flag = True # Assure in training mode model.train() crit = utils.LanguageModelCriterion() rl_crit = utils.RewardCriterion() optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay) # Load the optimizer if vars(opt).get('start_from', None) is not None and os.path.isfile(os.path.join(opt.start_from,"optimizer.pth")): optimizer.load_state_dict(torch.load(os.path.join(opt.start_from, 'optimizer.pth'))) while True: if update_lr_flag: # Assign the learning rate if epoch > opt.learning_rate_decay_start and opt.learning_rate_decay_start >= 0: frac = (epoch - opt.learning_rate_decay_start) // opt.learning_rate_decay_every decay_factor = opt.learning_rate_decay_rate ** frac opt.current_lr = opt.learning_rate * decay_factor else: opt.current_lr = opt.learning_rate utils.set_lr(optimizer, opt.current_lr) # Assign the scheduled sampling prob if epoch > opt.scheduled_sampling_start and opt.scheduled_sampling_start >= 0: frac = (epoch - opt.scheduled_sampling_start) // opt.scheduled_sampling_increase_every opt.ss_prob = min(opt.scheduled_sampling_increase_prob * frac, opt.scheduled_sampling_max_prob) model.ss_prob = opt.ss_prob # If start self critical training if opt.self_critical_after != -1 and epoch >= opt.self_critical_after: sc_flag = True init_cider_scorer(opt.cached_tokens) else: sc_flag = False update_lr_flag = False start = time.time() # Load data from train split (0) data = loader.get_batch('train') print('Read data:', time.time() - start) torch.cuda.synchronize() start = time.time() tmp = [data['fc_feats'], data['att_feats'], data['labels'], data['masks']] tmp = [Variable(torch.from_numpy(_), requires_grad=False).cuda() for _ in tmp] fc_feats, att_feats, labels, masks = tmp optimizer.zero_grad() if not sc_flag: loss = crit(model(fc_feats, att_feats, labels), labels[:,1:], masks[:,1:]) else: gen_result, sample_logprobs = model.sample(fc_feats, att_feats, {'sample_max':0}) reward = get_self_critical_reward(model, fc_feats, att_feats, data, gen_result) loss = rl_crit(sample_logprobs, gen_result, Variable(torch.from_numpy(reward).float().cuda(), requires_grad=False)) loss.backward() utils.clip_gradient(optimizer, opt.grad_clip) optimizer.step() train_loss = loss.data[0] torch.cuda.synchronize() end = time.time() if not sc_flag: print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \ .format(iteration, epoch, train_loss, end - start)) else: print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \ .format(iteration, epoch, np.mean(reward[:,0]), end - start)) # Update the iteration and epoch iteration += 1 if data['bounds']['wrapped']: epoch += 1 update_lr_flag = True # Write the training loss summary if (iteration % opt.losses_log_every == 0): if tf is not None: add_summary_value(tf_summary_writer, 'train_loss', train_loss, iteration) add_summary_value(tf_summary_writer, 'learning_rate', opt.current_lr, iteration) add_summary_value(tf_summary_writer, 'scheduled_sampling_prob', model.ss_prob, iteration) if sc_flag: add_summary_value(tf_summary_writer, 'avg_reward', np.mean(reward[:,0]), iteration) tf_summary_writer.flush() loss_history[iteration] = train_loss if not sc_flag else np.mean(reward[:,0]) lr_history[iteration] = opt.current_lr ss_prob_history[iteration] = model.ss_prob # make evaluation on validation set, and save model if (iteration % opt.save_checkpoint_every == 0): # eval model eval_kwargs = {'split': 'val', 'dataset': opt.input_json} eval_kwargs.update(vars(opt)) val_loss, predictions, lang_stats = eval_utils.eval_split(model, crit, loader, eval_kwargs) # Write validation result into summary if tf is not None: add_summary_value(tf_summary_writer, 'validation loss', val_loss, iteration) if lang_stats is not None: for k,v in lang_stats.items(): add_summary_value(tf_summary_writer, k, v, iteration) tf_summary_writer.flush() val_result_history[iteration] = {'loss': val_loss, 'lang_stats': lang_stats, 'predictions': predictions} # Save model if is improving on validation result if opt.language_eval == 1: current_score = lang_stats['CIDEr'] else: current_score = - val_loss best_flag = False if True: # if true if best_val_score is None or current_score > best_val_score: best_val_score = current_score best_flag = True checkpoint_path = os.path.join(opt.checkpoint_path, 'model.pth') torch.save(model.state_dict(), checkpoint_path) print("model saved to {}".format(checkpoint_path)) optimizer_path = os.path.join(opt.checkpoint_path, 'optimizer.pth') torch.save(optimizer.state_dict(), optimizer_path) # Dump miscalleous informations infos['iter'] = iteration infos['epoch'] = epoch infos['iterators'] = loader.iterators infos['split_ix'] = loader.split_ix infos['best_val_score'] = best_val_score infos['opt'] = opt infos['vocab'] = loader.get_vocab() histories['val_result_history'] = val_result_history histories['loss_history'] = loss_history histories['lr_history'] = lr_history histories['ss_prob_history'] = ss_prob_history with open(os.path.join(opt.checkpoint_path, 'infos_'+opt.id+'.pkl'), 'wb') as f: cPickle.dump(infos, f) with open(os.path.join(opt.checkpoint_path, 'histories_'+opt.id+'.pkl'), 'wb') as f: cPickle.dump(histories, f) if best_flag: checkpoint_path = os.path.join(opt.checkpoint_path, 'model-best.pth') torch.save(model.state_dict(), checkpoint_path) print("model saved to {}".format(checkpoint_path)) with open(os.path.join(opt.checkpoint_path, 'infos_'+opt.id+'-best.pkl'), 'wb') as f: cPickle.dump(infos, f) # Stop if reaching max epochs if epoch >= opt.max_epochs and opt.max_epochs != -1: break
def train(loader, model, optimizer, lr_scheduler, opt, device, crit): # Add Tensorboard train_logger, valid_logger = None, None if opt["log_dir"] is not None: train_logger = tb.SummaryWriter(path.join(opt["log_dir"], 'train'), flush_secs=1) # valid_logger = tb.SummaryWriter(path.join(opt["log_dir"], 'valid'), flush_secs=1) if opt["model"] == 'S2VTACTModel': use_action = True else: use_action = False # Training Procedure model.train() global_step = 0 # model = nn.DataParallel(model) # just ignore data parallel here for epoch in range(opt["epochs"]): iteration = 0 # If start self crit training if opt["self_crit_after"] != -1 and epoch >= opt["self_crit_after"]: sc_flag = True init_cider_scorer(opt["cached_tokens"]) else: sc_flag = False for data in loader: if torch.cuda.is_available(): torch.cuda.synchronize() fc_feats = data['fc_feats'].to(device) labels = data['labels'].to(device) masks = data['masks'].to(device) if use_action: action = data['action'].to(device) optimizer.zero_grad() if not sc_flag: if use_action: seq_probs, _ = model(vid_feats=fc_feats, action=action, device=device, target_variable=labels, mode='train') else: seq_probs, _ = model(vid_feats=fc_feats, target_variable=labels, mode='train') # Using Language Model Loss (NLLLoss or CrossEntropy Loss) loss = crit(seq_probs, labels[:, 1:], masks[:, 1:]) else: print('Currently ignore RL criterion') # seq_probs, seq_preds = model( # fc_feats, mode='inference', opt=opt) # reward = get_self_critical_reward(model, fc_feats, data, # seq_preds) # print(reward.shape) # loss = rl_crit(seq_probs, seq_preds, # torch.from_numpy(reward).float().cuda()) loss.backward() clip_grad_value_(model.parameters(), opt['grad_clip']) optimizer.step() train_loss = loss.item() if torch.cuda.is_available(): torch.cuda.synchronize() iteration += 1 global_step += 1 if not sc_flag: if iteration % 20 == 0: print("iter %d (epoch %d), train_loss = %.6f" % (iteration, epoch, train_loss)) else: print('Currently ignore RL criterion') # print("iter %d (epoch %d), avg_reward = %.6f" % # (iteration, epoch, np.mean(reward[:, 0]))) # Add Logger if train_logger is not None and global_step % 100 == 0: # Log some real data pass # Add Loss Statistics if train_logger is not None and iteration % 10 == 0: train_logger.add_scalar('loss', train_loss, global_step=global_step) # Step the Learning Rate Scheduler lr_scheduler.step() train_logger.add_scalar('lr', optimizer.param_groups[0]['lr'], global_step=global_step) if epoch % opt["save_checkpoint_every"] == 0: model_path = os.path.join(opt["checkpoint_path"], 'model_%d.pth' % (epoch)) model_info_path = os.path.join(opt["checkpoint_path"], 'model_score.txt') torch.save(model.state_dict(), model_path) print("model saved to %s" % (model_path)) with open(model_info_path, 'a') as f: f.write("model_%d, loss: %.6f\n" % (epoch, train_loss))
def train(loader, model, crit, optimizer, lr_scheduler, opt, rl_crit=None, opt_test=None, test_dataset=None): model.train() loss_avg = averager() #model = nn.DataParallel(model) writer = SummaryWriter() for epoch in range(opt["epochs"]): lr_scheduler.step() iteration = 0 # If start self crit training if opt["self_crit_after"] != -1 and epoch >= opt["self_crit_after"]: sc_flag = True init_cider_scorer(opt["cached_tokens"]) else: sc_flag = False for data in loader: torch.cuda.synchronize() fc_feats = data['fc_feats'].cuda() labels = data['labels'].cuda() masks = data['masks'].cuda() # clip_nums = data['clip_num'] # sorted_clip_nums, indices = torch.sort(clip_nums, descending=True) # _, desorted_indices = torch.sort(indices, descending=False) # fc_feats = fc_feats[indices] # pack = rnn.pack_padded_sequence(fc_feats, sorted_clip_nums, batch_first=True) optimizer.zero_grad() if not sc_flag: seq_probs, _ = model(fc_feats, labels, 'train') loss = crit(seq_probs, labels[:, 1:], masks[:, 1:]) else: seq_probs, seq_preds = model(fc_feats, mode='inference', opt=opt) reward = get_self_critical_reward(model, fc_feats, data, seq_preds) print(reward.shape) loss = rl_crit(seq_probs, seq_preds, torch.from_numpy(reward).float().cuda()) loss_avg.add(loss) loss.backward() clip_grad_value_(model.parameters(), opt['grad_clip']) optimizer.step() # train_loss = loss.item() torch.cuda.synchronize() iteration += 1 # if not sc_flag: # print("iter %d (epoch %d), train_loss = %.6f" % # (iteration, epoch, train_loss)) # else: # print("iter %d (epoch %d), avg_reward = %.6f" % # (iteration, epoch, np.mean(reward[:, 0]))) print("[epoch %d]->train_loss = %.6f" % (epoch, loss_avg.val())) writer.add_scalar('scalar/train_loss_epcho', loss_avg.val()) if epoch % opt["save_checkpoint_every"] == 0: test(model, crit, test_dataset, test_dataset.get_vocab(), opt_test, writer) model.train() model_path = os.path.join(opt["checkpoint_path"], 'model_%d.pth' % (epoch)) model_info_path = os.path.join(opt["checkpoint_path"], 'model_score.txt') torch.save(model.state_dict(), model_path) print("model saved to %s" % (model_path)) with open(model_info_path, 'a') as f: f.write("model_%d, loss: %.6f\n" % (epoch, loss_avg.val())) loss_avg.reset()