def init(opt):
infos = {}
histories = {}
if opt.seed > 0: torch.manual_seed(opt.seed)
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box:
opt.att_feat_size = opt.att_feat_size + 5 # box means [prob, x,y,w,h]
loader = DataLoader(opt) # Load image captioning dataset
coco_loader = DataLoader_COCO(opt) if opt.coco_eval_flag else None
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
if opt.start_from is not None and len(opt.start_from) > 0:
print('Start from: {}, Infos: {}'.format(
opt.start_from, os.path.join(opt.start_from, 'infos-best.pkl')))
# open old infos and check if models are compatible
with open(os.path.join(opt.start_from, 'infos-best.pkl')) as f:
infos = cPickle.load(f)
saved_model_opt = infos['opt']
need_be_same = ["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-best.pkl')):
with open(os.path.join(opt.start_from, 'histories-best.pkl')) as f:
histories = cPickle.load(f)
return opt, loader, coco_loader, infos, histories
def main():
import opts
import misc.utils as utils
opt = opts.parse_opt()
opt.caption_model ='topdown'
opt.batch_size=10
opt.id ='topdown'
opt.learning_rate= 5e-4
opt.learning_rate_decay_start= 0
opt.scheduled_sampling_start=0
opt.save_checkpoint_every=25#11500
opt.val_images_use=5000
opt.max_epochs=40
opt.start_from=None
opt.input_json='data/meta_coco_en.json'
opt.input_label_h5='data/label_coco_en.h5'
opt.input_image_h5 = 'data/coco_image_512.h5'
opt.use_att = utils.if_use_att(opt.caption_model)
opt.ccg = False
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
data = loader.get_batch('train')
data = loader.get_batch('val')
def show(opt):
opt.use_att = utils.if_use_att(opt.caption_model)
loader = DataLoader(opt, is_show=True)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
infos = {}
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
model = models.setup(opt)
model.cuda()
model.load_state_dict(torch.load(os.path.join(opt.start_from, 'model-best.pth')))
crit = utils.LanguageModelCriterion()
# eval model
eval_kwargs = {}
eval_kwargs.update(vars(opt))
eval_kwargs.update({'split': 'show',
'dataset': opt.input_json,
'language_eval': 0,
'beam_size': 5,
'print_all_beam': True})
val_loss, predictions, lang_stats = eval_utils.eval_split(model, crit, loader, eval_kwargs)
def show(opt):
opt.use_att = utils.if_use_att(opt.caption_model)
loader = DataLoader(opt, is_show=True)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
infos = {}
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
model = models.setup(opt)
model.cuda()
model.load_state_dict(
torch.load(os.path.join(opt.start_from, 'model-best.pth')))
crit = utils.LanguageModelCriterion()
# eval model
eval_kwargs = {}
eval_kwargs.update(vars(opt))
eval_kwargs.update({
'split': 'show',
'dataset': opt.input_json,
'language_eval': 0,
'beam_size': 5,
'print_all_beam': True
})
val_loss, predictions, lang_stats = eval_utils.eval_split(
model, crit, loader, eval_kwargs)
def main(args):
opt = vars(args)
# initialize
opt['dataset_splitBy'] = opt['dataset'] + '_' + opt['splitBy']
checkpoint_dir = osp.join(opt['checkpoint_path'], opt['dataset_splitBy'])
if not osp.isdir(checkpoint_dir): os.makedirs(checkpoint_dir)
# set random seed
torch.manual_seed(opt['seed'])
random.seed(opt['seed'])
# set up loader
data_json = osp.join('cache/prepro', opt['dataset_splitBy'], 'data.json')
data_h5 = osp.join('cache/prepro', opt['dataset_splitBy'], 'data.h5')
loader = CycleLoader(data_json, data_h5) ####
# set up model
opt['vocab_size']= loader.vocab_size
#opt['fc7_dim'] = loader.fc7_dim
#opt['pool5_dim'] = loader.pool5_dim
#opt['num_atts'] = loader.num_atts
#model = JointMatching(opt)
opt['C4_feat_dim'] = 1024
opt['use_att'] = utils.if_use_att(opt['caption_model'])
opt['seq_length'] = loader.label_length
#### can change to restore opt from info.pkl
#infos = {}
#histories = {}
if opt['start_from']is not None:
# open old infos and check if models are compatible
with open(os.path.join(opt['dataset_splitBy'], opt['start_from'], 'infos-best.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] == opt[checkme], "Command line argument and saved model disagree on '%s'" % checkme
#if os.path.isfile(os.path.join(opt['dataset_splitBy'], opt['start_from'], 'histories.pkl')):
# with open(os.path.join(opt['dataset_splitBy'], opt['start_from'], 'histories.pkl')) as f:
# histories = cPickle.load(f)
net = resnetv1(opt, batch_size=opt['batch_size'], num_layers=101) #### determine batch size in opt.py
# output directory where the models are saved
output_dir = osp.join(opt['dataset_splitBy'], 'output_{}'.format(opt['output_postfix']))
print('Output will be saved to `{:s}`'.format(output_dir))
# tensorboard directory where the summaries are saved during training
tb_dir = osp.join(opt['dataset_splitBy'], 'tb_{}'.format(opt['output_postfix']))
print('TensorFlow summaries will be saved to `{:s}`'.format(tb_dir))
# also add the validation set, but with no flipping images
orgflip = cfg.TRAIN.USE_FLIPPED
cfg.TRAIN.USE_FLIPPED = False
#_, valroidb = combined_roidb('coco_2014_minival')
#_, valroidb = combined_roidb('refcoco_test')
#print('{:d} validation roidb entries'.format(len(valroidb)))
cfg.TRAIN.USE_FLIPPED = orgflip
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
#train_net(net, imdb, roidb, valroidb, output_dir, tb_dir,
train_net(net, loader, output_dir, tb_dir,
pretrained_model='pyutils/mask-faster-rcnn/output/res101/coco_2014_train_minus_refer_valtest+coco_2014_valminusminival/notime/res101_mask_rcnn_iter_1250000.pth',
max_iters=args.max_iters)
def train(opt):
"""
:param caption decoder
:param VSE model : image encoder + caption encoder
"""
""" loading VSE model """
# Construct the model
vse = VSE(opt)
opt.best = os.path.join('./vse/model_best.pth.tar')
print("=> loading best checkpoint '{}'".format(opt.best))
checkpoint = torch.load(opt.best)
vse.load_state_dict(checkpoint['model'])
vse.val_start()
""" loading caption model """
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
save_path = os.path.join(opt.checkpoint_path,'CSGD')
if not os.path.exists(save_path):
os.makedirs(save_path, 0777)
infos = {}
histories = {}
RL_trainmodel = os.path.join('RL_%s' % opt.caption_model)
if opt.start_from is not None:
# open old infos and check if models are compatible
start_from_path = os.path.join(opt.start_from,'CSGD')
with open(os.path.join(start_from_path,'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(start_from_path, 'histories_'+opt.id+'.pkl')):
with open(os.path.join(start_from_path, 'histories_'+opt.id+'.pkl')) as f:
histories = cPickle.load(f)
with open(os.path.join(RL_trainmodel,'MLE','infos_'+opt.id+'-best.pkl')) as f:
infos_XE = cPickle.load(f)
opt.learning_rate = infos_XE['opt'].current_lr
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)
print(loader.iterators)
if opt.load_best_score == 1:
best_val_score = infos.get('best_val_score', None)
model = models.setup_pro(opt)
if vars(opt).get('start_from', None) is not None:
start_from_path = os.path.join(opt.start_from,'CSGD')
# check if all necessary files exist
assert os.path.isdir(opt.start_from)," %s must be a path" % opt.start_from
assert os.path.isfile(os.path.join(start_from_path,"infos_"+opt.id+".pkl")),"infos.pkl file does not exist in path %s"%opt.start_from
assert os.path.isfile(os.path.join(start_from_path,"optimizer.pth")) ,"optimizer.pth.file does not exist in path %s"%opt.start_from
model_path = os.path.join(start_from_path,'model.pth')
optimizer_path = os.path.join(start_from_path,'optimizer.pth')
else:
model_path = os.path.join(RL_trainmodel,'MLE', 'model-best.pth')
optimizer_path = os.path.join(RL_trainmodel,'MLE','optimizer-best.pth')
model.load_state_dict(torch.load(model_path))
print("model load from {}".format(model_path))
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)
optimizer.load_state_dict(torch.load(optimizer_path))
print("optimizer load from {}".format(optimizer_path))
all_cider = 0 # for computing the average CIDEr score
all_dis = 0 # for computing the discriminability percentage
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# Assign the scheduled sampling prob
model.ss_prob = 0.25
print('learning_rate: %s' %str(opt.current_lr))
update_lr_flag = False
# start self critical training
sc_flag = True
data = loader.get_batch('train')
torch.cuda.synchronize()
start = time.time()
# forward the model to also get generated samples for each image
# Only leave one feature for each image, in case duplicate sample
tmp = [data['fc_feats'][np.arange(loader.batch_size) * loader.seq_per_img],
data['att_feats'][np.arange(loader.batch_size) * loader.seq_per_img],
data['knn_fc_feats'][np.arange(loader.batch_size) * loader.seq_per_img],
data['knn_att_feats'][np.arange(loader.batch_size) * loader.seq_per_img]]
tmp = [Variable(torch.from_numpy(_), requires_grad=False).cuda() for _ in tmp]
fc_feats, att_feats, knn_fc_feats, knn_att_feats = tmp
optimizer.zero_grad()
gen_result, sample_logprobs = model.sample_score(fc_feats, att_feats, loader, {'sample_max': 0})
gen_result_baseline, sample_b_logprobs = model.sample_score(fc_feats, att_feats, loader, {'sample_max': 0})
bd_reward, sample_loss = get_bd_reward(vse, model, fc_feats, att_feats, data, gen_result,gen_result_baseline, loader)
hd_reward = get_hd_reward(vse, model, fc_feats, knn_fc_feats, data, gen_result,gen_result_baseline, loader)
cs_reward, m_cider = get_cs_reward(model, fc_feats, att_feats, data, gen_result, gen_result_baseline, loader)
reward = cs_reward - opt.hdr_w * hd_reward - opt.bdr_w * bd_reward
loss = rl_crit(sample_logprobs, gen_result, Variable(torch.from_numpy(reward).float().cuda(), requires_grad=False))
dis_number = (sample_loss < 0.4).float()
dis_number = dis_number.data.cpu().numpy().sum()
all_dis += dis_number
all_cider += m_cider
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), hdr = {:.3f}, bdr = {:.3f}, csr = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(hd_reward[:,0]), np.mean(bd_reward[:,0]), np.mean(cs_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):
loss_history[iteration] = 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 = evalpro_utils.eval_split(model, crit, loader, eval_kwargs)
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
save_path1 = os.path.join(save_path, 'model.pth')
if not os.path.exists(os.path.dirname(save_path1)):
os.makedirs(os.path.dirname(save_path1))
torch.save(model.state_dict(), save_path1)
print("model saved to {}".format(save_path1))
optimizer_path1 = os.path.join(save_path, 'optimizer.pth')
if not os.path.exists(os.path.dirname(optimizer_path1)):
os.makedirs(os.path.dirname(optimizer_path1))
torch.save(optimizer.state_dict(), optimizer_path1)
print("optimizer saved to {}".format(optimizer_path1))
all_dis = all_dis / opt.save_checkpoint_every
print("all_dis:%f" %all_dis)
infos['all_dis'] = all_dis
all_cider = all_cider / opt.save_checkpoint_every
print("all_cider:%f" %all_cider)
infos['all_cider'] = all_cider
# 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(save_path, 'infos_'+opt.id+'.pkl'), 'wb') as f:
cPickle.dump(infos, f)
with open(os.path.join(save_path, 'histories_'+opt.id+'.pkl'), 'wb') as f:
cPickle.dump(histories, f)
if best_flag:
save_path2 = os.path.join(save_path, 'model-best.pth')
torch.save(model.state_dict(), save_path2)
optimizer_path2 = os.path.join(save_path, 'optimizer-best.pth')
torch.save(optimizer.state_dict(), optimizer_path2)
print("model saved to {}".format(save_path2))
with open(os.path.join(save_path, 'infos_'+opt.id+'-best.pkl'), 'wb') as f:
cPickle.dump(infos, f)
with open(os.path.join(save_path,'histories_'+opt.id+'-best.pkl'), 'wb') as f:
cPickle.dump(histories, f)
# Stop if reaching max epochs
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
def train(opt):
if vars(opt).get('start_from', None) is not None:
opt.checkpoint_path = opt.start_from
opt.id = opt.checkpoint_path.split('/')[-1]
print('Point to folder: {}'.format(opt.checkpoint_path))
else:
opt.id = datetime.datetime.now().strftime(
'%Y%m%d_%H%M%S') + '_' + opt.caption_model
opt.checkpoint_path = os.path.join(opt.checkpoint_path, opt.id)
if not os.path.exists(opt.checkpoint_path):
os.makedirs(opt.checkpoint_path)
print('Create folder: {}'.format(opt.checkpoint_path))
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
# opt.use_att = False
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
loader = DataLoader_UP(opt)
opt.vocab_size = loader.vocab_size
if opt.use_rela == 1:
opt.rela_dict_size = loader.rela_dict_size
opt.seq_length = loader.seq_length
use_rela = getattr(opt, 'use_rela', 0)
try:
tb_summary_writer = tf and tf.compat.v1.summary.FileWriter(
opt.checkpoint_path)
except:
print('Set tensorboard error!')
pdb.set_trace()
infos = {}
histories = {}
if opt.start_from is not None:
# open old infos and check if models are compatible
with open(os.path.join(opt.checkpoint_path, 'infos.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.checkpoint_path, 'histories.pkl')):
with open(os.path.join(opt.checkpoint_path, 'histories.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).cuda()
# dp_model = torch.nn.DataParallel(model)
# dp_model = torch.nn.DataParallel(model, [0,2,3])
dp_model = model
print('### Model summary below###\n {}\n'.format(str(model)))
model_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('model parameter:{}'.format(model_params))
update_lr_flag = True
# Assure in training mode
dp_model.train()
parameters = model.named_children()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = utils.build_optimizer(
filter(lambda p: p.requires_grad, model.parameters()), opt)
optimizer.zero_grad()
accumulate_iter = 0
train_loss = 0
reward = np.zeros([1, 1])
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_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(opt.train_split)
# print('Read data:', time.time() - start)
torch.cuda.synchronize()
start = time.time()
fc_feats = None
att_feats = None
att_masks = None
ssg_data = None
rela_data = None
if getattr(opt, 'use_ssg', 0) == 1:
if getattr(opt, 'use_isg', 0) == 1:
tmp = [
data['fc_feats'], data['labels'], data['masks'],
data['att_feats'], data['att_masks'],
data['isg_rela_matrix'], data['isg_rela_masks'],
data['isg_obj'], data['isg_obj_masks'], data['isg_attr'],
data['isg_attr_masks'], data['ssg_rela_matrix'],
data['ssg_rela_masks'], data['ssg_obj'],
data['ssg_obj_masks'], data['ssg_attr'],
data['ssg_attr_masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, labels, masks, att_feats, att_masks, \
isg_rela_matrix, isg_rela_masks, isg_obj, isg_obj_masks, isg_attr, isg_attr_masks, \
ssg_rela_matrix, ssg_rela_masks, ssg_obj, ssg_obj_masks, ssg_attr, ssg_attr_masks = tmp
# image graph domain
isg_data = {}
isg_data['att_feats'] = att_feats
isg_data['att_masks'] = att_masks
isg_data['isg_rela_matrix'] = isg_rela_matrix
isg_data['isg_rela_masks'] = isg_rela_masks
isg_data['isg_obj'] = isg_obj
isg_data['isg_obj_masks'] = isg_obj_masks
isg_data['isg_attr'] = isg_attr
isg_data['isg_attr_masks'] = isg_attr_masks
# text graph domain
ssg_data = {}
ssg_data['ssg_rela_matrix'] = ssg_rela_matrix
ssg_data['ssg_rela_masks'] = ssg_rela_masks
ssg_data['ssg_obj'] = ssg_obj
ssg_data['ssg_obj_masks'] = ssg_obj_masks
ssg_data['ssg_attr'] = ssg_attr
ssg_data['ssg_attr_masks'] = ssg_attr_masks
else:
tmp = [
data['fc_feats'], data['ssg_rela_matrix'],
data['ssg_rela_masks'], data['ssg_obj'],
data['ssg_obj_masks'], data['ssg_attr'],
data['ssg_attr_masks'], data['labels'], data['masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, ssg_rela_matrix, ssg_rela_masks, ssg_obj, ssg_obj_masks, ssg_attr, ssg_attr_masks, labels, masks = tmp
ssg_data = {}
ssg_data['ssg_rela_matrix'] = ssg_rela_matrix
ssg_data['ssg_rela_masks'] = ssg_rela_masks
ssg_data['ssg_obj'] = ssg_obj
ssg_data['ssg_obj_masks'] = ssg_obj_masks
ssg_data['ssg_attr'] = ssg_attr
isg_data = None
ssg_data['ssg_attr_masks'] = ssg_attr_masks
else:
tmp = [
data['fc_feats'], data['att_feats'], data['labels'],
data['masks'], data['att_masks']
]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, masks, att_masks = tmp
if not sc_flag:
# loss = crit(dp_model(model_zh,model_en,itow_zh,itow, fc_feats, labels, isg_data, ssg_data), labels[:, 1:], masks[:, 1:])
loss = crit(dp_model(fc_feats, labels, isg_data, ssg_data),
labels[:, 1:], masks[:, 1:])
else:
gen_result, sample_logprobs = dp_model(fc_feats,
isg_data,
ssg_data,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats, isg_data,
ssg_data, data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
accumulate_iter = accumulate_iter + 1
loss = loss / opt.accumulate_number
loss.backward()
if accumulate_iter % opt.accumulate_number == 0:
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
optimizer.zero_grad()
iteration += 1
accumulate_iter = 0
train_loss = loss.item() * opt.accumulate_number
end = time.time()
if not sc_flag:
print("{}/{}/{}|train_loss={:.3f}|time/batch={:.3f}" \
.format(opt.id, iteration, epoch, train_loss, end - start))
else:
print("{}/{}/{}|avg_reward={:.3f}|time/batch={:.3f}" \
.format(opt.id, iteration, epoch, np.mean(reward[:, 0]), end - start))
torch.cuda.synchronize()
# Update the iteration and epoch
if data['bounds']['wrapped']:
epoch += 1
update_lr_flag = True
# Write the training loss summary
if (iteration % opt.losses_log_every == 0) and (iteration != 0):
add_summary_value(tb_summary_writer, 'train_loss', train_loss,
iteration)
add_summary_value(tb_summary_writer, 'learning_rate',
opt.current_lr, iteration)
add_summary_value(tb_summary_writer, 'scheduled_sampling_prob',
model.ss_prob, iteration)
if sc_flag:
add_summary_value(tb_summary_writer, 'avg_reward',
np.mean(reward[:, 0]), iteration)
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 %10 == 0) and (iteration != 0):
if (iteration % opt.save_checkpoint_every == 0) and (iteration != 0):
# eval model
if use_rela:
eval_kwargs = {
'split': 'val',
'dataset': opt.input_json,
'use_real': 1
}
else:
eval_kwargs = {'split': 'val', 'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
# val_loss, predictions, lang_stats = eval_utils.eval_split(model_zh,model_en,itow_zh,itow, dp_model, crit, loader, eval_kwargs)
val_loss, predictions, lang_stats = eval_utils.eval_split(
dp_model, crit, loader, eval_kwargs)
# Write validation result into summary
add_summary_value(tb_summary_writer, 'validation loss', val_loss,
iteration)
if lang_stats is not None:
for k, v in lang_stats.items():
add_summary_value(tb_summary_writer, k, v, iteration)
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
save_id = iteration / opt.save_checkpoint_every
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.pkl'),
'wb') as f:
cPickle.dump(infos, f)
with open(os.path.join(opt.checkpoint_path, 'histories.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-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(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
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()
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:
optimizer.load_state_dict(
torch.load(os.path.join(opt.start_from, 'optimizer.pth')))
train_loss_list = []
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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
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()
loss = crit(model(fc_feats, att_feats, labels), labels[:, 1:],
masks[:, 1:])
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, end - start))
train_loss_list.append(train_loss)
# 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):
loss_history[iteration] = train_loss
lr_history[iteration] = opt.current_lr
ss_prob_history[iteration] = model.ss_prob
# make evaluation on validation set, and save model
min_val_loss = 100
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, temp_min = eval_utils.eval_split(
model, crit, loader, eval_kwargs)
if (temp_min < min_val_loss):
min_val_loss = temp_min
# Write validation result into summary
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:
# eval model
eval_kwargs = {'split': 'val', 'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats, temp_min = eval_utils.eval_split(
model, crit, loader, eval_kwargs)
if (temp_min < min_val_loss):
min_val_loss = temp_min
# Write validation result into summary
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)
plt.title('Training loss')
plt.xlabel('iteration')
plt.ylabel('loss')
plt.plot(np.arange(1, len(train_loss_list) + 1), train_loss_list)
savefilename = 'lab3-1.jpg'
plt.savefig(savefilename)
plt.close()
print("min loss is :", min_val_loss)
break
import time
import os
from six.moves import cPickle
import opts
import models
from dataloader import *
import eval_utils
import misc.utils as utils
from misc.rewards import init_scorer, get_self_critical_reward
opt = opts.parse_opt()
# Setting up model
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
print(opt.checkpoint_path)
infos = {}
histories = {}
if opt.start_from is not None:
# open old infos and check if models are compatible
print(os.getcwd())
with open(
os.path.join(os.getcwd(), opt.start_from,
def train(opt):
import random
random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(0)
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
from dataloader_pair import DataLoader
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
tb_summary_writer = tb and tb.SummaryWriter(opt.checkpoint_path)
if opt.log_to_file:
if os.path.exists(os.path.join(opt.checkpoint_path, 'log')):
suffix = time.strftime("%Y-%m-%d %X", time.localtime())
print('Warning !!! %s already exists ! use suffix ! ' %
os.path.join(opt.checkpoint_path, 'log'))
sys.stdout = open(
os.path.join(opt.checkpoint_path, 'log' + suffix), "w")
else:
print('logging to file %s' %
os.path.join(opt.checkpoint_path, 'log'))
sys.stdout = open(os.path.join(opt.checkpoint_path, 'log'), "w")
infos = {}
histories = {}
if opt.start_from is not None:
# open old infos and check if models are compatible
if os.path.isfile(opt.start_from):
with open(os.path.join(opt.infos)) 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
else:
if opt.load_best != 0:
print('loading best info')
with open(
os.path.join(opt.start_from,
'infos_' + opt.id + '-best.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
else:
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'), 'rb') as f:
try:
histories = cPickle.load(f)
except:
print('load history error!')
histories = {}
iteration = infos.get('iter', 0)
epoch = infos.get('epoch', 0)
start_epoch = epoch
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).cuda()
dp_model = torch.nn.DataParallel(model)
update_lr_flag = True
# Assure in training mode
dp_model.train()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = utils.build_optimizer(model.parameters(), opt)
#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')))
if opt.caption_model == 'att2in2p':
optimized = [
'logit2', 'ctx2att2', 'core2', 'prev_sent_emb', 'prev_sent_wrap'
]
optimized_param = []
optimized_param1 = []
for name, param in model.named_parameters():
second = False
for n in optimized:
if n in name:
print('second', name)
optimized_param.append(param)
second = True
if 'embed' in name:
print('all', name)
optimized_param1.append(param)
optimized_param.append(param)
elif not second:
print('first', name)
optimized_param1.append(param)
while True:
if opt.val_only:
eval_kwargs = {'split': 'val', 'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
print('start evaluating')
val_loss, predictions, lang_stats = eval_utils_pair.eval_split(
dp_model, crit, loader, eval_kwargs)
exit(0)
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_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['pair_fc_feats'], data['pair_att_feats'], data['pair_labels'],
data['pair_masks'], data['pair_att_masks']
]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, masks, att_masks = tmp
masks = masks.float()
optimizer.zero_grad()
if not sc_flag:
if opt.onlysecond:
# only using the second sentence from a visual paraphrase pair. opt.caption_model should be a one-stage decoding model
loss = crit(
dp_model(fc_feats, att_feats, labels[:, 1, :], att_masks),
labels[:, 1, 1:], masks[:, 1, 1:])
loss1 = loss2 = loss / 2
elif opt.first:
# using the first sentence
tmp = [
data['first_fc_feats'], data['first_att_feats'],
data['first_labels'], data['first_masks'],
data['first_att_masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, att_feats, labels, masks, att_masks = tmp
masks = masks.float()
loss = crit(
dp_model(fc_feats, att_feats, labels[:, :], att_masks),
labels[:, 1:], masks[:, 1:])
loss1 = loss2 = loss / 2
elif opt.onlyfirst:
# only using the second sentence from a visual paraphrase pair
loss = crit(
dp_model(fc_feats, att_feats, labels[:, 0, :], att_masks),
labels[:, 0, 1:], masks[:, 0, 1:])
loss1 = loss2 = loss / 2
else:
# proposed DCVP model, opt.caption_model should be att2inp
output1, output2 = dp_model(fc_feats, att_feats, labels,
att_masks, masks[:, 0, 1:])
loss1 = crit(output1, labels[:, 0, 1:], masks[:, 0, 1:])
loss2 = crit(output2, labels[:, 1, 1:], masks[:, 1, 1:])
loss = loss1 + loss2
else:
raise NotImplementedError
# Our DCVP model does not support self-critical sequence training
# We found that RL(SCST) with CIDEr reward will improve conventional metrics (BLEU, CIDEr, etc.)
# but harm diversity and descriptiveness
# Please refer to the paper for the details
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.item()
torch.cuda.synchronize()
end = time.time()
if not sc_flag:
print("iter {} (epoch {}), train_loss = {:.3f}, loss1 = {:.3f}, loss2 = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, loss.item(), loss1.item(), loss2.item(), end - start))
else:
print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(reward[:,0]), end - start))
sys.stdout.flush()
# 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):
add_summary_value(tb_summary_writer, 'train_loss', train_loss,
iteration)
add_summary_value(tb_summary_writer, 'learning_rate',
opt.current_lr, iteration)
add_summary_value(tb_summary_writer, 'scheduled_sampling_prob',
model.ss_prob, iteration)
if sc_flag:
add_summary_value(tb_summary_writer, 'avg_reward',
np.mean(reward[:, 0]), iteration)
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_pair.eval_split(
dp_model, crit, loader, eval_kwargs)
# Write validation result into summary
add_summary_value(tb_summary_writer, 'validation loss', val_loss,
iteration)
if lang_stats is not None:
for k, v in lang_stats.items():
add_summary_value(tb_summary_writer, k, v, iteration)
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)
checkpoint_path = os.path.join(
opt.checkpoint_path, 'model' + str(iteration) + '.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 + '_' + str(iteration) + '.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(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
def ids_to_sents(ids):
return utils.decode_sequence(loader, ids)
tf_summary_writer = tf and tf.summary.FileWriter(opt.checkpoint_path)
def load_infos(dir=opt.start_from, suffix=''):
# open old infos and check if models are compatible
with open(os.path.join(dir, 'infos_{}{}.pkl'.format(opt.id,
suffix))) 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 getattr(saved_model_opt, checkme) == getattr(
opt, checkme
), "Command line argument and saved model disagree on '%s'" % checkme
return infos
def load_histories(dir=opt.start_from, suffix=''):
path = os.path.join(dir, 'histories_{}{}.pkl'.format(opt.id, suffix))
if os.path.isfile(path):
with open(path) as f:
histories = cPickle.load(f)
return histories
infos = {}
histories = {}
if opt.start_from is not None:
infos = load_infos()
histories = load_histories()
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_ce = utils.LanguageModelCriterion()
crit_mb = mBLEU(4)
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:
optimizer.load_state_dict(
torch.load(os.path.join(opt.start_from, 'optimizer.pth')))
def eval_model():
model.eval()
eval_kwargs = {'split': 'val', 'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats = eval_utils.eval_split(
model, crit_ce, loader, eval_kwargs)
# Write validation result into summary
if tf is not None:
add_summary_value(tf_summary_writer, 'validation loss', val_loss,
iteration)
for k, v in lang_stats.items():
add_summary_value(tf_summary_writer, k, v, iteration)
tf_summary_writer.flush()
model.train()
return val_loss, predictions, lang_stats
eval_model()
opt.current_teach_mask_prefix_length = opt.teach_mask_prefix_length
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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
# Assign the teach mask prefix length
if epoch > opt.teach_mask_prefix_length_increase_start:
frac = (epoch - opt.teach_mask_prefix_length_increase_start
) // opt.teach_mask_prefix_length_increase_every
opt.current_teach_mask_prefix_length = opt.teach_mask_prefix_length + frac * opt.teach_mask_prefix_length_increase_steps
update_lr_flag = False
verbose = (iteration % opt.verbose_iters == 0)
start = time.time()
# Load data from train split (0)
data = loader.get_batch('train')
if iteration % opt.print_iters == 0:
print('Read data:', time.time() - start)
torch.cuda.synchronize()
start = time.time()
tmp = [
data['fc_feats'], data['att_feats'], data['labels'], data['masks']
]
tmp = [torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, masks = tmp
optimizer.zero_grad()
teach_mask = utils.make_teach_mask(labels.size(1), opt)
enable_ce = (opt.bleu_w != 1)
enable_mb = (opt.bleu_w != 0)
if enable_ce:
enable_xe = (opt.xe_w != 0)
enable_pg = (opt.pg_w != 0)
if enable_xe:
logits = model(
fc_feats,
att_feats,
labels,
teach_mask=(teach_mask if opt.teach_ce
and not opt.teach_all_input else None))
if opt.teach_ce:
decode_length = logits.shape[1] + 1
teach_mask = teach_mask[:decode_length]
onehot = utils.to_onehot(labels[:, :decode_length],
logits.shape[-1],
dtype=torch.float)
probs = torch.exp(logits)
probs = torch.cat([onehot[:, :1], probs], 1)
probs = utils.mask_probs(probs, onehot, teach_mask)
if verbose:
verbose_probs = probs
verbose_probs.retain_grad()
logits = torch.log(1. - (1. - 1e-6) * (1. - probs))[:, 1:]
loss_xe = crit_ce(logits, labels[:, 1:], masks[:, 1:])
else:
loss_xe = 0.
if enable_pg:
ids_sample, logprobs_sample = model.sample(
fc_feats, att_feats, opt={'sample_max': 0})
ids_greedy, logprobs_greedy = model.sample(
fc_feats, att_feats, opt={'sample_max': 1})
seq_sample = utils.tolist(ids_sample)
seq_greedy = utils.tolist(ids_greedy)
seq_target = utils.tolist(labels[:, 1:])
rewards = [
sentence_bleu([t], s, smooth=True) -
sentence_bleu([t], g, smooth=True)
for s, g, t in zip(seq_sample, seq_greedy, seq_target)
]
rewards = torch.tensor(rewards, device='cuda')
mask_sample = torch.ne(ids_sample,
torch.tensor(0, device='cuda')).float()
loss_pg = (rewards *
(logprobs_sample * mask_sample).sum(1)).mean()
else:
loss_pg = 0.
loss_ce = opt.xe_w * loss_xe + opt.pg_w * loss_pg
else:
loss_ce = 0.
if enable_mb:
logits = model(
fc_feats,
att_feats,
labels,
teach_mask=(teach_mask if not opt.teach_all_input else None))
decode_length = logits.shape[1] + 1
teach_mask = teach_mask[:decode_length]
onehot = utils.to_onehot(labels[:, :decode_length],
logits.shape[-1],
dtype=torch.float)
probs = torch.exp(logits)
probs = torch.cat([onehot[:, :1], probs], 1) # pad bos
probs = utils.mask_probs(probs, onehot, teach_mask)
if verbose:
verbose_probs = probs
verbose_probs.retain_grad()
mask = masks[:, :decode_length]
mask = torch.cat([mask[:, :1], mask], 1)
loss_mb = crit_mb(probs,
labels[:, :decode_length],
mask,
min_fn=opt.min_fn,
min_c=opt.min_c,
verbose=verbose)
else:
loss_mb = 0.
loss = loss_ce * (1 - opt.bleu_w) + loss_mb * opt.bleu_w
loss.backward()
utils.clip_gradient(
optimizer,
opt.grad_clip) #TODO: examine clip method and record grad
if verbose and 'verbose_probs' in locals():
max_grads, max_ids = verbose_probs.grad.topk(opt.verbose_topk,
-1,
largest=False)
max_probs = torch.gather(verbose_probs, -1, max_ids)
max_sents = ids_to_sents(max_ids[:, :, 0])
for sample_i in range(min(opt.samples, verbose_probs.shape[0])):
l = len(max_sents[sample_i]) + 1
print('max:\n{}'.format(max_sents[sample_i]))
print('max probs:\n{}'.format(max_probs[sample_i][:l]))
print('max grads:\n{}'.format(max_grads[sample_i][:l]))
optimizer.step()
train_loss = float(loss)
torch.cuda.synchronize()
end = time.time()
if iteration % opt.print_iters == 0:
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, 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)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
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):
val_loss, predictions, lang_stats = eval_model()
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
if True: # if true
best_flag = False
if best_val_score is None or current_score > best_val_score:
best_val_score = current_score
best_flag = True
# 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
histories['val_result_history'] = val_result_history
histories['loss_history'] = loss_history
histories['lr_history'] = lr_history
histories['ss_prob_history'] = ss_prob_history
def save_model(suffix=''):
model_path = os.path.join(opt.checkpoint_path,
'model{}.pth'.format(suffix))
torch.save(model.state_dict(), model_path)
print("model saved to {}".format(model_path))
optimizer_path = os.path.join(
opt.checkpoint_path, 'optimizer{}.pth'.format(suffix))
torch.save(optimizer.state_dict(), optimizer_path)
with open(
os.path.join(
opt.checkpoint_path,
'infos_{}{}.pkl'.format(opt.id, suffix)),
'wb') as f:
cPickle.dump(infos, f)
with open(
os.path.join(
opt.checkpoint_path,
'histories_{}{}.pkl'.format(opt.id, suffix)),
'wb') as f:
cPickle.dump(histories, f)
save_model()
save_model(".iter{}".format(iteration))
if best_flag:
save_model(".best")
# Stop if reaching max epochs
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
def train(opt):
opt.use_att = utils.if_use_att(opt.caption_model)
from dataloader import DataLoader
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.vocab_ccg_size = loader.vocab_ccg_size
opt.seq_length = loader.seq_length
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)
cnn_model = utils.build_cnn(opt)
cnn_model.cuda()
model = models.setup(opt)
model.cuda()
# model = DataParallel(model)
if vars(opt).get('start_from', None) is not None:
# check if all necessary files exist
assert os.path.isdir(
opt.start_from), " %s must be a a path" % opt.start_from
assert os.path.isfile(
os.path.join(opt.start_from, "infos_" + opt.id + ".pkl")
), "infos.pkl file does not exist in path %s" % opt.start_from
model.load_state_dict(
torch.load(os.path.join(opt.start_from, 'model.pth')))
update_lr_flag = True
model.train()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
multilabel_crit = nn.MultiLabelSoftMarginLoss().cuda()
# optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay)
optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate)
if opt.finetune_cnn_after != -1 and epoch >= opt.finetune_cnn_after:
print('finetune mode')
cnn_optimizer = optim.Adam([\
{'params': module.parameters()} for module in cnn_model._modules.values()[5:]\
], lr=opt.cnn_learning_rate, weight_decay=opt.cnn_weight_decay)
if vars(opt).get('start_from', None) is not None and os.path.isfile(
os.path.join(opt.start_from, "optimizer.pth")):
if 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')))
if opt.finetune_cnn_after != -1 and epoch >= opt.finetune_cnn_after:
if os.path.isfile(os.path.join(opt.start_from,
'optimizer-cnn.pth')):
cnn_optimizer.load_state_dict(
torch.load(
os.path.join(opt.start_from, 'optimizer-cnn.pth')))
eval_kwargs = {'split': 'val', 'dataset': opt.input_json, 'verbose': True}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats = eval_utils.eval_split(
cnn_model, model, crit, loader, eval_kwargs, True)
epoch_start = time.time()
while True:
if update_lr_flag:
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
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
#model.module.ss_prob = opt.ss_prob
if opt.self_critical_after != -1 and epoch >= opt.self_critical_after:
sc_flag = True
else:
sc_flag = False
# Update the training stage of cnn
for p in cnn_model.parameters():
p.requires_grad = True
# Fix the first few layers:
for module in cnn_model._modules.values()[:5]:
for p in module.parameters():
p.requires_grad = False
cnn_model.train()
update_lr_flag = False
cnn_model.apply(utils.set_bn_fix)
cnn_model.apply(utils.set_bn_eval)
start = time.time()
torch.cuda.synchronize()
data = loader.get_batch('train')
if opt.finetune_cnn_after != -1 and epoch >= opt.finetune_cnn_after:
multilabels = [
data['detection_infos'][i]['label']
for i in range(len(data['detection_infos']))
]
tmp = [
data['labels'], data['masks'],
np.array(multilabels, dtype=np.int16)
]
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
labels, masks, multilabels = tmp
images = data[
'images'] # it cannot be turned into tensor since different sizes.
_fc_feats_2048 = []
_fc_feats_81 = []
_att_feats = []
for i in range(loader.batch_size):
x = Variable(torch.from_numpy(images[i]),
requires_grad=False).cuda()
x = x.unsqueeze(0)
att_feats, fc_feats_81 = cnn_model(x)
fc_feats_2048 = att_feats.mean(3).mean(2).squeeze()
att_feats = F.adaptive_avg_pool2d(att_feats,
[14, 14]).squeeze().permute(
1, 2, 0) #(0, 2, 3, 1)
_fc_feats_2048.append(fc_feats_2048)
_fc_feats_81.append(fc_feats_81)
_att_feats.append(att_feats)
_fc_feats_2048 = torch.stack(_fc_feats_2048)
_fc_feats_81 = torch.stack(_fc_feats_81)
_att_feats = torch.stack(_att_feats)
att_feats = _att_feats.unsqueeze(1).expand(*((_att_feats.size(0), loader.seq_per_img,) + \
_att_feats.size()[1:])).contiguous().view(*((_att_feats.size(0) * loader.seq_per_img,) + \
_att_feats.size()[1:]))
fc_feats_2048 = _fc_feats_2048.unsqueeze(1).expand(*((_fc_feats_2048.size(0), loader.seq_per_img,) + \
_fc_feats_2048.size()[1:])).contiguous().view(*((_fc_feats_2048.size(0) * loader.seq_per_img,) + \
_fc_feats_2048.size()[1:]))
fc_feats_81 = _fc_feats_81
#
cnn_optimizer.zero_grad()
else:
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:
loss1 = crit(model(fc_feats_2048, att_feats, labels),
labels[:, 1:], masks[:, 1:])
loss2 = multilabel_crit(fc_feats_81.double(), multilabels.double())
loss = 0.8 * loss1 + 0.2 * loss2.float()
else:
gen_result, sample_logprobs = model.sample(fc_feats_2048,
att_feats,
{'sample_max': 0})
reward = get_self_critical_reward(model, fc_feats_2048, att_feats,
data, gen_result)
loss1 = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(reward).float().cuda(),
requires_grad=False))
loss2 = multilabel_crit(fc_feats_81.double(), multilabels.double())
loss3 = crit(model(fc_feats_2048, att_feats, labels),
labels[:, 1:], masks[:, 1:])
loss = 0.995 * loss1 + 0.005 * (loss2.float() + loss3)
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
mle_loss = loss1.data[0]
multilabel_loss = loss2.data[0]
torch.cuda.synchronize()
end = time.time()
if not sc_flag and iteration % 2500 == 0:
print("iter {} (epoch {}), mle_loss = {:.3f}, multilabel_loss = {:.3f}, train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, mle_loss, multilabel_loss, train_loss, end - start))
if sc_flag and iteration % 2500 == 0:
print("iter {} (epoch {}), avg_reward = {:.3f}, mle_loss = {:.3f}, multilabel_loss = {:.3f}, train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(reward[:,0]), mle_loss, multilabel_loss, train_loss, end - start))
iteration += 1
if (iteration % opt.losses_log_every == 0):
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
if (iteration % opt.save_checkpoint_every == 0):
eval_kwargs = {
'split': 'val',
'dataset': opt.input_json,
'verbose': True
}
eval_kwargs.update(vars(opt))
if opt.finetune_cnn_after != -1 and epoch >= opt.finetune_cnn_after:
val_loss, predictions, lang_stats = eval_utils.eval_split(
cnn_model, model, crit, loader, eval_kwargs, True)
else:
val_loss, predictions, lang_stats = eval_utils.eval_split(
cnn_model, model, crit, loader, eval_kwargs, False)
val_result_history[iteration] = {
'loss': val_loss,
'lang_stats': lang_stats,
'predictions': predictions
}
if opt.language_eval == 1:
current_score = lang_stats['CIDEr']
else:
current_score = -val_loss
best_flag = False
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))
cnn_checkpoint_path = os.path.join(opt.checkpoint_path,
'model-cnn.pth')
torch.save(cnn_model.state_dict(), cnn_checkpoint_path)
print("cnn model saved to {}".format(cnn_checkpoint_path))
optimizer_path = os.path.join(opt.checkpoint_path,
'optimizer.pth')
torch.save(optimizer.state_dict(), optimizer_path)
if opt.finetune_cnn_after != -1 and epoch >= opt.finetune_cnn_after:
cnn_optimizer_path = os.path.join(opt.checkpoint_path,
'optimizer-cnn.pth')
torch.save(cnn_optimizer.state_dict(), cnn_optimizer_path)
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))
cnn_checkpoint_path = os.path.join(opt.checkpoint_path,
'model-cnn-best.pth')
torch.save(cnn_model.state_dict(), cnn_checkpoint_path)
print("cnn model saved to {}".format(cnn_checkpoint_path))
with open(
os.path.join(opt.checkpoint_path,
'infos_' + opt.id + '-best.pkl'),
'wb') as f:
cPickle.dump(infos, f)
if data['bounds']['wrapped']:
epoch += 1
update_lr_flag = True
print("epoch: " + str(epoch) + " during: " +
str(time.time() - epoch_start))
epoch_start = time.time()
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
def main():
import opts
import misc.utils as utils
opt = opts.parse_opt()
opt.caption_model = 'topdown'
opt.batch_size = 10 #512#32*4*4
opt.id = 'topdown'
opt.learning_rate = 5e-4
opt.learning_rate_decay_start = 0
opt.scheduled_sampling_start = 0
opt.save_checkpoint_every = 5000 #450#5000#11500
opt.val_images_use = 5000
opt.max_epochs = 50 #30
opt.start_from = 'save/rt' #"save" #None
opt.language_eval = 1
opt.input_json = 'data/meta_coco_en.json'
opt.input_label_h5 = 'data/label_coco_en.h5'
# opt.input_json='data/coco_ccg.json' #'data/meta_coco_en.json'
# opt.input_label_h5='data/coco_ccg_label.h5' #'data/label_coco_en.h5'
# opt.input_fc_dir='/nlp/andyweizhao/self-critical.pytorch-master/data/cocotalk_fc'
# opt.input_att_dir='/nlp/andyweizhao/self-critical.pytorch-master/data/cocotalk_att'
opt.finetune_cnn_after = 0
opt.ccg = False
opt.input_image_h5 = 'data/coco_image_512.h5'
opt.use_att = utils.if_use_att(opt.caption_model)
from dataloader import DataLoader # just-in-time generated features
loader = DataLoader(opt)
# from dataloader_fixcnn import DataLoader # load pre-processed features
# loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.vocab_ccg_size = loader.vocab_ccg_size
opt.seq_length = loader.seq_length
import models
model = models.setup(opt)
cnn_model = utils.build_cnn(opt)
cnn_model.cuda()
model.cuda()
data = loader.get_batch('train')
images = data['images']
# _fc_feats_2048 = []
# _fc_feats_81 = []
# _att_feats = []
# for i in range(loader.batch_size):
# x = Variable(torch.from_numpy(images[i]), volatile=True).cuda()
# x = x.unsqueeze(0)
# att_feats, fc_feats_81 = cnn_model(x)
# fc_feats_2048 = att_feats.mean(3).mean(2).squeeze()
# att_feats = F.adaptive_avg_pool2d(att_feats,[14,14]).squeeze().permute(1, 2, 0)#(0, 2, 3, 1)
# _fc_feats_2048.append(fc_feats_2048)
# _fc_feats_81.append(fc_feats_81)
# _att_feats.append(att_feats)
# _fc_feats_2048 = torch.stack(_fc_feats_2048)
# _fc_feats_81 = torch.stack(_fc_feats_81)
# _att_feats = torch.stack(_att_feats)
# att_feats = _att_feats.unsqueeze(1).expand(*((_att_feats.size(0), loader.seq_per_img,) + \
# _att_feats.size()[1:])).contiguous().view(*((_att_feats.size(0) * loader.seq_per_img,) + \
# _att_feats.size()[1:]))
# fc_feats_2048 = _fc_feats_2048.unsqueeze(1).expand(*((_fc_feats_2048.size(0), loader.seq_per_img,) + \
# _fc_feats_2048.size()[1:])).contiguous().view(*((_fc_feats_2048.size(0) * loader.seq_per_img,) + \
# _fc_feats_2048.size()[1:]))
# fc_feats_81 = _fc_feats_81
#
# att_feats = Variable(att_feats, requires_grad=False).cuda()
# Variable(fc_feats_81)
crit = utils.LanguageModelCriterion()
eval_kwargs = {'split': 'val', 'dataset': opt.input_json, 'verbose': True}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats = eval_split(cnn_model, model, crit,
loader, eval_kwargs, True)
def train(opt):
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
tb_summary_writer = tb and tb.SummaryWriter(opt.checkpoint_path)
infos = {}
histories = {}
if opt.start_from_path is not None:
# open old infos and check if models are compatible
with open(os.path.join(opt.start_from_path,
'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_path,
'histories_' + opt.id + '.pkl')):
with open(
os.path.join(opt.start_from_path,
'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', {})
#print(val_result_history.get(3000))
#exit(0)
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).cuda()
no = sum(p.numel() for p in model.parameters())
pytorch_total_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("Trainable Params:" + str(pytorch_total_params))
print("Total Params:" + str(no))
#exit(0)
dp_model = torch.nn.DataParallel(model)
epoch_done = True
# Assure in training mode
dp_model.train()
if (opt.use_obj_mcl_loss == 1):
mcl_crit = utils.MultiLabelClassification()
if opt.label_smoothing > 0:
crit = utils.LabelSmoothing(smoothing=opt.label_smoothing)
else:
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
if opt.noamopt:
assert opt.caption_model == 'transformer', 'noamopt can only work with transformer'
optimizer = utils.get_std_opt(model,
factor=opt.noamopt_factor,
warmup=opt.noamopt_warmup)
optimizer._step = iteration
elif opt.reduce_on_plateau:
optimizer = utils.build_optimizer(model.parameters(), opt)
optimizer = utils.ReduceLROnPlateau(optimizer, factor=0.5, patience=3)
else:
optimizer = utils.build_optimizer(model.parameters(), opt)
# Load the optimizer
if vars(opt).get('start_from_path', None) is not None and os.path.isfile(
os.path.join(opt.start_from_path, "optimizer.pth")):
optimizer.load_state_dict(
torch.load(os.path.join(opt.start_from_path, 'optimizer.pth')))
time_epoch_start = time.time()
data_time_sum = 0
batch_time_sum = 0
while True:
if epoch_done:
torch.cuda.synchronize()
time_epoch_end = time.time()
time_elapsed = (time_epoch_end - time_epoch_start)
print('[DEBUG] Epoch Time: ' + str(time_elapsed))
print('[DEBUG] Sum Data Time: ' + str(data_time_sum))
print('[DEBUG] Sum Batch Time: ' + str(batch_time_sum))
#if epoch==1:
# exit(0)
if not opt.noamopt and not opt.reduce_on_plateau:
# 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) # set the decayed rate
# 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_scorer(opt.cached_tokens)
else:
sc_flag = False
epoch_done = False
start = time.time()
# Load data from train split (0)
data = loader.get_batch('train')
print('Read data:', time.time() - start)
data_time_sum += time.time() - start
torch.cuda.synchronize()
start = time.time()
if (opt.use_obj_mcl_loss == 0):
tmp = [
data['fc_feats'], data['att_feats'], data['labels'],
data['masks'], data['att_masks']
]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, masks, att_masks = tmp
else:
if opt.use_obj_att and opt.use_seg_feat:
tmp = [
data['fc_feats'], data['att_feats'], data['obj_att_feats'],
data['seg_feat_feats'], data['labels'], data['masks'],
data['obj_labels'], data['att_masks'],
data['obj_att_masks'], data['seg_feat_masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, att_feats, obj_att_feats, seg_feat_feats, labels, masks, obj_labels, att_masks, obj_att_masks, seg_feat_masks = tmp
elif not opt.use_obj_att and opt.use_seg_feat:
tmp = [
data['fc_feats'], data['att_feats'],
data['seg_feat_feats'], data['labels'], data['masks'],
data['obj_labels'], data['att_masks'],
data['seg_feat_masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, att_feats, seg_feat_feats, labels, masks, obj_labels, att_masks, seg_feat_masks = tmp
elif not opt.use_obj_att and not opt.use_seg_feat:
tmp = [
data['fc_feats'], data['att_feats'], data['labels'],
data['masks'], data['obj_labels'], data['att_masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, att_feats, labels, masks, obj_labels, att_masks = tmp
elif opt.use_obj_att and not opt.use_seg_feat:
tmp = [
data['fc_feats'], data['att_feats'], data['obj_att_feats'],
data['labels'], data['masks'], data['obj_labels'],
data['att_masks'], data['obj_att_masks']
]
tmp = [
_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp
]
fc_feats, att_feats, obj_att_feats, labels, masks, obj_labels, att_masks, obj_att_masks = tmp
optimizer.zero_grad()
if (opt.use_obj_mcl_loss == 0):
if not sc_flag:
loss = crit(dp_model(fc_feats, att_feats, labels, att_masks),
labels[:, 1:], masks[:, 1:])
else:
gen_result, sample_logprobs = dp_model(fc_feats,
att_feats,
att_masks,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats,
att_feats, att_masks, data,
gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
else:
if opt.use_obj_att and opt.use_seg_feat:
if not sc_flag:
logits, out = dp_model(
fc_feats, [att_feats, obj_att_feats, seg_feat_feats],
labels, [att_masks, obj_att_masks, seg_feat_masks])
caption_loss = crit(logits, labels[:, 1:], masks[:, 1:])
obj_loss = mcl_crit(out, obj_labels)
loss = opt.lambda_caption * caption_loss + opt.lambda_obj * obj_loss
#loss = 0.1*caption_loss + obj_loss
#loss = caption_loss + 0 * obj_loss
else:
gen_result, sample_logprobs = dp_model(
fc_feats,
att_feats,
att_masks,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats,
att_feats, att_masks,
data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
elif not opt.use_obj_att and opt.use_seg_feat:
if not sc_flag:
logits, out = dp_model(fc_feats,
[att_feats, seg_feat_feats], labels,
[att_masks, seg_feat_masks])
caption_loss = crit(logits, labels[:, 1:], masks[:, 1:])
obj_loss = mcl_crit(out, obj_labels)
loss = opt.lambda_caption * caption_loss + opt.lambda_obj * obj_loss
#loss = caption_loss + 0 * obj_loss
else:
gen_result, sample_logprobs = dp_model(
fc_feats,
att_feats,
att_masks,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats,
att_feats, att_masks,
data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
if not opt.use_obj_att and not opt.use_seg_feat:
if not sc_flag:
logits, out = dp_model(fc_feats, att_feats, labels,
att_masks)
caption_loss = crit(logits, labels[:, 1:], masks[:, 1:])
obj_loss = mcl_crit(out, obj_labels)
loss = opt.lambda_caption * caption_loss + opt.lambda_obj * obj_loss
#loss = caption_loss + 0 * obj_loss
else:
gen_result, sample_logprobs = dp_model(
fc_feats,
att_feats,
att_masks,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats,
att_feats, att_masks,
data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
elif opt.use_obj_att and not opt.use_seg_feat:
if not sc_flag:
logits, out = dp_model(fc_feats,
[att_feats, obj_att_feats], labels,
[att_masks, obj_att_masks])
caption_loss = crit(logits, labels[:, 1:], masks[:, 1:])
obj_loss = mcl_crit(out, obj_labels)
loss = 0.1 * caption_loss + obj_loss
#loss = caption_loss + 0 * obj_loss
else:
gen_result, sample_logprobs = dp_model(
fc_feats,
att_feats,
att_masks,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats,
att_feats, att_masks,
data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.item()
torch.cuda.synchronize()
end = time.time()
batch_time_sum += end - start
if not sc_flag:
if (opt.use_obj_mcl_loss == 1):
print("iter {} (epoch {}), train_loss = {:.3f}, caption_loss = {:.3f}, object_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, caption_loss.item(), obj_loss.item(), end - start))
else:
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
epoch_done = True
# Write the training loss summary
if (iteration % opt.losses_log_every == 0):
add_summary_value(tb_summary_writer, 'train_loss', train_loss,
iteration)
if (opt.use_obj_mcl_loss == 1):
add_summary_value(tb_summary_writer, 'obj_loss',
obj_loss.item(), iteration)
add_summary_value(tb_summary_writer, 'caption_loss',
caption_loss.item(), iteration)
if opt.noamopt:
opt.current_lr = optimizer.rate()
elif opt.reduce_on_plateau:
opt.current_lr = optimizer.current_lr
add_summary_value(tb_summary_writer, 'learning_rate',
opt.current_lr, iteration)
add_summary_value(tb_summary_writer, 'scheduled_sampling_prob',
model.ss_prob, iteration)
if sc_flag:
add_summary_value(tb_summary_writer, 'avg_reward',
np.mean(reward[:, 0]), iteration)
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
orig_batch_size = opt.batch_size
opt.batch_size = 1
eval_kwargs = {'split': 'val', 'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
loader.batch_size = eval_kwargs.get('batch_size', 1)
val_loss, predictions, lang_stats = eval_utils.eval_split(
dp_model, crit, loader, eval_kwargs)
opt.batch_size = orig_batch_size
loader.batch_size = orig_batch_size
if opt.reduce_on_plateau:
if 'CIDEr' in lang_stats:
optimizer.scheduler_step(-lang_stats['CIDEr'])
else:
optimizer.scheduler_step(val_loss)
# Write validation result into summary
add_summary_value(tb_summary_writer, 'validation loss', val_loss,
iteration)
for k, v in lang_stats.items():
add_summary_value(tb_summary_writer, k, v, iteration)
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 eval_cap(opt):
opt.start_from = None
opt.start_from_gan = 'data/save_gan_rcsls_s_g_naacl/20201121_185728_sep_self_att_sep_gan_only'
opt.checkpoint_path_gan = opt.start_from_gan
opt.init_path_zh = 'data/save_for_finetune/20201116_203753_gtssg_sep_self_att_sep_vv1_RCSLS_submap_v2_add_wordmap_add_global_naacl_self_gate_p2/model-best.pth'
opt.input_isg_dir = "data/coco_graph_extract_ft_isg_joint_rcsls_submap_global_naacl_self_gate_finetune"
opt.input_ssg_dir = "data/coco_graph_extract_ft_ssg_joint_rcsls_submap_global_naacl_self_gate_finetune"
opt.caption_model_to_replace = 'up_gtssg_sep_self_att_sep_vv1_RCSLS_submap_v2_add_wordmap_add_global_naacl_self_gate_p2'
opt.gpu = 0
opt.batch_size = 50
opt.beam_size = 5
opt.dump_path = 1
opt.caption_model = 'sep_self_att_sep_gan_only'
opt.input_json = 'data/coco_cn/cocobu_gan_ssg.json'
opt.input_json_isg = 'data/coco_cn/cocobu_gan_isg.json'
opt.input_label_h5 = 'data/coco_cn/cocobu_gan_isg_label.h5'
opt.ssg_dict_path = 'data/aic_process/ALL_11683_v3_COCOCN_spice_sg_dict_t5.npz_revise.npz'
opt.rela_dict_dir = 'data/rela_dict.npy'
opt.input_fc_dir = 'data/cocobu_fc'
opt.input_att_dir = 'data/cocobu_att'
opt.input_box_dir = 'data/cocotalk_box'
opt.input_label_h5 = 'data/cocobu_label.h5'
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
opt.split = 'test'
loader = DataLoader_GAN(opt)
loader_i2t = DataLoader_UP(opt)
opt.vocab_size = loader.vocab_size
if opt.use_rela == 1:
opt.rela_dict_size = loader.rela_dict_size
opt.seq_length = loader.seq_length
infos = {}
histories = {}
if opt.start_from is not None:
# open old infos and check if models are compatible
try:
with open(os.path.join(opt.checkpoint_path, 'infos.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.checkpoint_path, 'histories.pkl')):
with open(os.path.join(opt.checkpoint_path,
'histories.pkl')) as f:
histories = cPickle.load(f)
except:
print("Can not load infos.pkl")
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)
# opt.caption_model = 'up_gtssg_sep_self_att_sep'
opt.caption_model = opt.caption_model_to_replace
model = models.setup(opt).cuda()
model.eval()
opt.start_from = opt.start_from_gan
opt.checkpoint_path = opt.start_from_gan
opt.id = opt.checkpoint_path.split('/')[-1]
start_from = opt.start_from_gan
with open(os.path.join(opt.checkpoint_path, 'infos.pkl')) as f:
infos = cPickle.load(f)
saved_model_opt = infos['opt']
saved_model_opt.start_from_gan = start_from
saved_model_opt.checkpoint_path_gan = start_from
netG_A_obj = GAN_init_G(saved_model_opt,
Generator(saved_model_opt),
type='netG_A_obj').cuda().eval()
netG_A_rel = GAN_init_G(saved_model_opt,
Generator(saved_model_opt),
type='netG_A_rel').cuda().eval()
netG_A_atr = GAN_init_G(saved_model_opt,
Generator(saved_model_opt),
type='netG_A_atr').cuda().eval()
netG_B_obj = GAN_init_G(saved_model_opt,
Generator(saved_model_opt),
type='netG_B_obj').cuda().eval()
netG_B_rel = GAN_init_G(saved_model_opt,
Generator(saved_model_opt),
type='netG_B_rel').cuda().eval()
netG_B_atr = GAN_init_G(saved_model_opt,
Generator(saved_model_opt),
type='netG_B_atr').cuda().eval()
val_loss, cache_path = eval_utils_gan.eval_split_gan_v2(
opt, model, netG_A_obj, netG_A_rel, netG_A_atr, netG_B_obj, netG_B_rel,
netG_B_atr, loader, loader_i2t, opt.split)
def train(opt):
assert opt.annfile is not None and len(opt.annfile) > 0
print('Checkpoint path is ' + opt.checkpoint_path)
print('This program is using GPU ' +
str(os.environ['CUDA_VISIBLE_DEVICES']))
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
tb_summary_writer = tb and tb.SummaryWriter(opt.checkpoint_path)
infos = {}
histories = {}
if opt.start_from is not None:
# open old infos and check if models are compatible
if opt.load_best:
info_path = os.path.join(opt.start_from,
'infos_' + opt.id + '-best.pkl')
else:
info_path = os.path.join(opt.start_from,
'infos_' + opt.id + '.pkl')
with open(info_path) 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)
if opt.learning_rate_decay_start is None:
opt.learning_rate_decay_start = infos.get(
'opt', None).learning_rate_decay_start
# if opt.load_best:
# opt.self_critical_after = epoch
elif opt.learning_rate_decay_start == -1 and opt.self_critical_after != -1 and epoch >= opt.self_critical_after:
opt.learning_rate_decay_start = epoch
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)
best_val_score_ave_model = infos.get('best_val_score_ave_model', None)
model = models.setup(opt).cuda()
dp_model = torch.nn.DataParallel(model)
update_lr_flag = True
# Assure in training mode
dp_model.train()
crit = utils.LanguageModelCriterion(opt.XE_eps)
rl_crit = utils.RewardCriterion()
# build_optimizer
optimizer = build_optimizer(model, opt)
# Load the optimizer
if opt.load_opti and vars(opt).get(
'start_from',
None) is not None and opt.load_best == 0 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')))
# initialize the running average of parameters
avg_param = deepcopy(list(p.data for p in model.parameters()))
# make evaluation using original model
best_val_score, histories, infos = eva_original_model(
best_val_score, crit, epoch, histories, infos, iteration, loader,
loss_history, lr_history, model, opt, optimizer, ss_prob_history,
tb_summary_writer, val_result_history)
while True:
if update_lr_flag:
# Assign the learning rate
if epoch > opt.learning_rate_decay_start and opt.learning_rate_decay_start >= 0:
if opt.lr_decay == 'exp':
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
elif opt.lr_decay == 'cosine':
lr_epoch = min((epoch - opt.learning_rate_decay_start),
opt.lr_max_epoch)
cosine_decay = 0.5 * (
1 + math.cos(math.pi * lr_epoch / opt.lr_max_epoch))
decay_factor = (1 - opt.lr_cosine_decay_base
) * cosine_decay + opt.lr_cosine_decay_base
opt.current_lr = opt.learning_rate * decay_factor
else:
opt.current_lr = opt.learning_rate
lr = [opt.current_lr]
if opt.att_normalize_method is not None and '6' in opt.att_normalize_method:
lr = [opt.current_lr, opt.lr_ratio * opt.current_lr]
utils.set_lr(optimizer, lr)
print('learning rate is: ' + str(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_scorer(opt.cached_tokens)
else:
sc_flag = False
update_lr_flag = False
# Update the iteration
iteration += 1
# Load data from train split (0)
data = loader.get_batch(opt.train_split)
torch.cuda.synchronize()
start = time.time()
tmp = [
data['fc_feats'], data['att_feats'], data['labels'], data['masks'],
data['att_masks']
]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, masks, att_masks = tmp
optimizer.zero_grad()
if not sc_flag:
output = dp_model(fc_feats, att_feats, labels, att_masks)
# calculate loss
loss = crit(output[0], labels[:, 1:], masks[:, 1:])
# add some middle variable histogram
if iteration % (4 * opt.losses_log_every) == 0:
outputs = [
_.data.cpu().numpy() if _ is not None else None
for _ in output
]
variables_histogram(data, iteration, outputs,
tb_summary_writer, opt)
else:
gen_result, sample_logprobs = dp_model(fc_feats,
att_feats,
att_masks,
opt={'sample_max': 0},
mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats, att_feats,
att_masks, data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
# grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_max_norm)
# add_summary_value(tb_summary_writer, 'grad_L2_norm', grad_norm, iteration)
optimizer.step()
train_loss = loss.item()
torch.cuda.synchronize()
end = time.time()
# compute the running average of parameters
for p, avg_p in zip(model.parameters(), avg_param):
avg_p.mul_(opt.beta).add_((1.0 - opt.beta), p.data)
if iteration % 10 == 0:
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 epoch
if data['bounds']['wrapped']:
epoch += 1
update_lr_flag = True
# Write the training loss summary
if (iteration % opt.losses_log_every == 0):
add_summary_value(tb_summary_writer, 'train_loss', train_loss,
iteration)
add_summary_value(tb_summary_writer, 'learning_rate',
opt.current_lr, iteration)
add_summary_value(tb_summary_writer, 'scheduled_sampling_prob',
model.ss_prob, iteration)
if sc_flag:
add_summary_value(tb_summary_writer, 'avg_reward',
np.mean(reward[:, 0]), iteration)
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
if opt.tensorboard_weights_grads and (iteration %
(8 * opt.losses_log_every) == 0):
# add weights histogram to tensorboard summary
for name, param in model.named_parameters():
if (opt.tensorboard_parameters_name is None or sum([
p_name in name
for p_name in opt.tensorboard_parameters_name
]) > 0) and param.grad is not None:
tb_summary_writer.add_histogram(
'Weights_' + name.replace('.', '/'), param, iteration)
tb_summary_writer.add_histogram(
'Grads_' + name.replace('.', '/'), param.grad,
iteration)
if opt.tensorboard_buffers and (iteration %
(opt.losses_log_every) == 0):
for name, buffer in model.named_buffers():
if (opt.tensorboard_buffers_name is None or sum([
p_name in name
for p_name in opt.tensorboard_buffers_name
]) > 0) and buffer is not None:
add_summary_value(tb_summary_writer,
name.replace('.',
'/'), buffer, iteration)
if opt.distance_sensitive_coefficient and iteration % (
4 * opt.losses_log_every) == 0:
print('The coefficient in intra_att_att_lstm is as follows:')
print(
model.core.intra_att_att_lstm.coefficient.data.cpu().tolist())
print('The coefficient in intra_att_lang_lstm is as follows:')
print(
model.core.intra_att_lang_lstm.coefficient.data.cpu().tolist())
if opt.distance_sensitive_bias and iteration % (
4 * opt.losses_log_every) == 0:
print('The bias in intra_att_att_lstm is as follows:')
print(model.core.intra_att_att_lstm.bias.data.cpu().tolist())
print('The bias in intra_att_lang_lstm is as follows:')
print(model.core.intra_att_lang_lstm.bias.data.cpu().tolist())
# make evaluation using original model
if (iteration % opt.save_checkpoint_every == 0):
best_val_score, histories, infos = eva_original_model(
best_val_score, crit, epoch, histories, infos, iteration,
loader, loss_history, lr_history, model, opt, optimizer,
ss_prob_history, tb_summary_writer, val_result_history)
# make evaluation with the averaged parameters model
if iteration > opt.ave_threshold and (iteration %
opt.save_checkpoint_every == 0):
best_val_score_ave_model, infos = eva_ave_model(
avg_param, best_val_score_ave_model, crit, infos, iteration,
loader, model, opt, tb_summary_writer)
# # Stop if reaching max epochs
# if epoch >= opt.max_epochs and opt.max_epochs != -1:
# break
if iteration >= opt.max_iter:
break
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_S is not None:
if os.path.isfile(
os.path.join(opt.start_from_S,
'histories_' + opt.id + '.pkl')):
with open(
os.path.join(opt.start_from_S,
'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)
# Set CommNetModel
model1 = ShowTellModel(opt)
model2 = ShowTellModel(opt)
model1.load_state_dict(
torch.load(os.path.join(opt.start_from_T, 'model.pth')))
model2.load_state_dict(
torch.load(os.path.join(opt.start_from_S, 'model.pth')))
model1.cuda()
model2.cuda()
model = CommNetModel(opt, model1, model2)
model.cuda()
logger = Logger(opt)
update_lr_flag = True
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_S', None) is not None and os.path.isfile(os.path.join(opt.start_from_S,"optimizer.pth")):
# temp = torch.load(os.path.join(opt.start_from_S, 'optimizer.pth'))
# optimizer.load_state_dict(temp)
while True:
if update_lr_flag:
opt, sc_flag, update_lr_flag, model, optimizer = update_lr(
opt, epoch, model, optimizer)
# Load data from train split (0)
data = loader.get_batch('train', seq_per_img=opt.seq_per_img)
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_1 = crit(
model(fc_feats, att_feats, labels)[0], labels[:, 1:],
masks[:, 1:])
loss_2 = crit(
model(fc_feats, att_feats, labels)[1], labels[:, 1:],
masks[:, 1:])
loss_1.backward()
loss_2.backward()
loss = loss_1 + loss_2
else:
gen_result_1, sample_logprobs_1 = model.model1.sample(
fc_feats, att_feats, {'sample_max': 0})
gen_result_2, sample_logprobs_2 = model.model2.sample(
fc_feats, att_feats, {'sample_max': 0})
reward_1 = get_self_critical_reward_forCommNet(
model, fc_feats, att_feats, data, gen_result_1, logger)
reward_2 = get_self_critical_reward_forCommNet(
model, fc_feats, att_feats, data, gen_result_2, logger)
loss_1 = rl_crit(
sample_logprobs_1, gen_result_1,
Variable(torch.from_numpy(reward_1).float().cuda(),
requires_grad=False))
loss_2 = rl_crit(
sample_logprobs_2, gen_result_2,
Variable(torch.from_numpy(reward_2).float().cuda(),
requires_grad=False))
loss_1.backward()
loss_2.backward()
loss = loss_1 + loss_2
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
if not sc_flag:
log = "iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, end - start)
logger.write(log)
else:
log = "iter {} (epoch {}), S_avg_reward = {:.3f}, T_avg_reward = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(reward_1[:,0]), np.mean(reward_2[:,0]), end - start)
logger.write(log)
# 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_S',
np.mean(reward_1[:, 0]), iteration)
add_summary_value(tf_summary_writer, 'avg_reward_T',
np.mean(reward_2[:, 0]), iteration)
tf_summary_writer.flush()
loss_history[iteration] = train_loss if not sc_flag else np.mean(
reward_1[:, 0] + reward_2[:, 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.model1, crit, loader, logger, eval_kwargs)
val_loss, predictions, lang_stats = eval_utils_forCommNet.eval_split(
model, crit, loader, logger, eval_kwargs)
logger.write_dict(lang_stats)
# Write validation result into summary
if tf is not None:
add_summary_value(tf_summary_writer, 'validation loss',
val_loss, iteration)
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,
'model1.pth')
torch.save(model.model1.state_dict(), checkpoint_path)
print("model1 saved to {}".format(checkpoint_path))
checkpoint_path = os.path.join(opt.checkpoint_path,
'model2.pth')
torch.save(model.model2.state_dict(), checkpoint_path)
print("model2 saved to {}".format(checkpoint_path))
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,
'model1-best.pth')
torch.save(model.model1.state_dict(), checkpoint_path)
print("model1 saved to {}".format(checkpoint_path))
checkpoint_path = os.path.join(opt.checkpoint_path,
'model2-best.pth')
torch.save(model.model2.state_dict(), checkpoint_path)
print("model2 saved to {}".format(checkpoint_path))
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(opt):
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
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).cuda()
dp_model = torch.nn.DataParallel(model)
update_lr_flag = True
# Assure in training mode
dp_model.train()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = utils.build_optimizer(model.parameters(), opt)
# 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_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')
data_time = time.time() - start
torch.cuda.synchronize()
start = time.time()
tmp = [data['fc_feats'], data['att_feats'], data['labels'], data['masks'], data['att_masks']]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, masks, att_masks = tmp
optimizer.zero_grad()
if not sc_flag:
loss = crit(dp_model(fc_feats, att_feats, labels, att_masks), labels[:,1:], masks[:,1:])
else:
gen_result, sample_logprobs = dp_model(fc_feats, att_feats, att_masks, opt={'sample_max':0}, mode='sample')
reward = get_self_critical_reward(dp_model, fc_feats, att_feats, att_masks, data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data, torch.from_numpy(reward).float().cuda())
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.item()
torch.cuda.synchronize()
end = time.time()
if iteration % opt.print_freq == 0:
if not sc_flag:
print("iter {} (epoch {}), train_loss = {:.3f}, batch time = {:.3f}, data time = {:.3f}" \
.format(iteration, epoch, train_loss, end - start, data_time))
else:
print("iter {} (epoch {}), avg_reward = {:.3f}, batch time = {:.3f}, data time = {:.3f}" \
.format(iteration, epoch, np.mean(reward[:,0]), end - start, data_time))
# 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):
checkpoint_path = os.path.join(opt.checkpoint_path, 'model.pth')
torch.save(model.state_dict(), checkpoint_path) # MODIFIED (ADDED)
# eval model
eval_kwargs = {'split': 'val',
'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats = eval_utils.eval_split(dp_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-i{}-score{}.pth'.format(iteration, best_val_score))
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(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'),'rb') 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'),'rb') 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()
if opt.multi_gpu:
model=nn.DataParallel(model)
update_lr_flag = True
# Assure in training mode
model.train()
crit = utils.LanguageModelCriterion()
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:
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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
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()
loss = crit(model(fc_feats, att_feats, labels), labels[:,1:], masks[:,1:])
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, 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)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
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)
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(opt):
start_from = vars(opt).get('start_from', None)
start_from_p = vars(opt).get('start_from_en', None)
opt.checkpoint_path, opt.id = model_start(start_from, 0)
opt.checkpoint_path_p, opt.id_p = model_start(start_from_p, 1)
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
# opt.use_att = False
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
loader = DataLoader_UP(opt)
vocab_size = loader.vocab_size
vocab_size_p = loader.vocab_size_p
if opt.use_rela == 1:
opt.rela_dict_size = loader.rela_dict_size
opt.seq_length = loader.seq_length
use_rela = getattr(opt, 'use_rela', 0)
try:
tb_summary_writer = tf and tf.compat.v1.summary.FileWriter(
opt.checkpoint_path)
tb_summary_writer_p = tf and tf.compat.v1.summary.FileWriter(
opt.checkpoint_path_p)
except:
print('Set tensorboard error!')
pdb.set_trace()
opt.p_flag = 0 # whether paired model
opt.vocab_size = vocab_size
loader,iteration,epoch,val_result_history,loss_history,lr_history,ss_prob_history,best_val_score,\
infos,histories,update_lr_flag,model,dp_model,parameters,crit,rl_crit,optimizer,accumulate_iter,train_loss,reward,train_loss_kl,train_loss_all=load_info(loader,start_from,opt.checkpoint_path,opt.p_flag)
opt.p_flag = 1
opt.vocab_size = vocab_size_p
loader,iteration_p,epoch_p,val_result_history_p,loss_history_p,lr_history_p,ss_prob_history_p,best_val_score_p, \
infos_p, histories_p, update_lr_flag_p, model_p,dp_model_p, parameters_p, crit_p, rl_crit_p, optimizer_p, accumulate_iter_p, train_loss_p, reward_p,train_loss_kl,train_loss_all = load_info(
loader, start_from_p,opt.checkpoint_path_p,opt.p_flag)
# # global variables
update_lr_flag = update_lr_flag
accumulate_iter = accumulate_iter
train_loss = train_loss
train_loss_kl = train_loss_kl
train_loss_all = train_loss_all
reward = reward
update_lr_flag_p = update_lr_flag_p
accumulate_iter_p = accumulate_iter_p
train_loss_p = train_loss_p
reward_p = reward_p
while True:
# Load data from train split (0)
data = loader.get_batch(opt.train_split)
# print('Read data:', time.time() - start)
torch.cuda.synchronize()
start = time.time()
opt.p_flag = 0 # whether paired model
loss, att_obj, att_rela, att_attr, update_lr_flag, sc_flag = pre_model(
update_lr_flag, epoch, optimizer, model, data, dp_model, crit,
rl_crit, opt.p_flag)
opt.p_flag = 1
loss_p, att_obj_p, att_rela_p, att_attr_p, update_lr_flag_p, sc_flag_p = pre_model(
update_lr_flag_p, epoch_p, optimizer_p, model_p, data, dp_model_p,
crit_p, rl_crit_p, opt.p_flag)
att_obj = F.softmax(att_obj, dim=1)
att_obj_p = F.softmax(att_obj_p, dim=1)
att_rela = F.softmax(att_rela, dim=1)
att_rela_p = F.softmax(att_rela_p, dim=1)
att_attr = F.softmax(att_attr, dim=1)
att_attr_p = F.softmax(att_attr_p, dim=1)
loss_kl = torch.exp(F.kl_div(att_obj.log(), att_obj_p, reduction='sum'),out=None)\
+torch.exp(F.kl_div(att_rela.log(), att_rela_p, reduction='sum'),out=None)\
+torch.exp(F.kl_div(att_attr.log(), att_attr_p, reduction='sum'),out=None)
# print(loss_kl)
accumulate_iter = accumulate_iter + 1
accumulate_iter_p = accumulate_iter_p + 1
loss_all = loss + loss_p + loss_kl
loss = loss / opt.accumulate_number
loss_p = loss_p / opt.accumulate_number
loss_kl = loss_kl / opt.accumulate_number
loss_all = loss_all / opt.accumulate_number
loss_all.backward()
opt.p_flag = 0
# print ('iteration of model 1 is {}'.format(iteration))
update_lr_flag, epoch, optimizer, model, dp_model, accumulate_iter, iteration, loss, sc_flag, start, reward, \
tb_summary_writer, loss_history, lr_history, ss_prob_history, use_rela, val_result_history, best_val_score, loader,\
infos, histories, train_loss, train_loss_kl, train_loss_all, loss_all, loss_kl=\
save_model(opt.input_json,accumulate_iter, optimizer, iteration, loss, sc_flag, epoch, start, reward, data, tb_summary_writer,model,
loss_history, lr_history, ss_prob_history, use_rela, dp_model, val_result_history, best_val_score,
crit, loader, infos, histories,train_loss, train_loss_kl ,train_loss_all,opt.id,opt.checkpoint_path,opt.p_flag,loss_all,loss_kl,update_lr_flag)
opt.p_flag = 1
# print('iteration of model 2 is {}'.format(iteration_p))
update_lr_flag_p, epoch_p, optimizer_p, model_p, dp_model_p, accumulate_iter_p, iteration_p, loss_p, sc_flag_p, start, reward_p, \
tb_summary_writer_p, loss_history_p, lr_history_p, ss_prob_history_p, use_rela, val_result_history_p, best_val_score_p, loader,\
infos_p, histories_p, train_loss_p, train_loss_kl, train_loss_all, loss_all, loss_kl=\
save_model(opt.input_json_en,accumulate_iter_p, optimizer_p, iteration_p, loss_p, sc_flag_p, epoch_p, start, reward_p, data, tb_summary_writer_p, model_p,
loss_history_p, lr_history_p, ss_prob_history_p, use_rela, dp_model_p, val_result_history_p, best_val_score_p,
crit_p, loader, infos_p, histories_p,train_loss_p,train_loss_kl ,train_loss_all,opt.id_p,opt.checkpoint_path_p,opt.p_flag,loss_all,loss_kl,update_lr_flag_p)
# Stop if reaching max epochs
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
def train(opt):
if vars(opt).get('start_from', None) is not None:
opt.checkpoint_path = opt.start_from
opt.id = opt.checkpoint_path.split('/')[-1]
print('Point to folder: {}'.format(opt.checkpoint_path))
else:
opt.id = datetime.datetime.now().strftime(
'%Y%m%d_%H%M%S') + '_' + opt.caption_model
opt.checkpoint_path = os.path.join(opt.checkpoint_path, opt.id)
if not os.path.exists(opt.checkpoint_path):
os.makedirs(opt.checkpoint_path)
print('Create folder: {}'.format(opt.checkpoint_path))
# Write YAML file
with io.open(opt.checkpoint_path + '/opts.yaml', 'w',
encoding='utf8') as outfile:
yaml.dump(opt, outfile, default_flow_style=False, allow_unicode=True)
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
loader = DataLoader_GAN(opt)
loader_i2t = DataLoader_UP(opt)
opt.vocab_size = loader.vocab_size
if opt.use_rela == 1:
opt.rela_dict_size = loader.rela_dict_size
opt.seq_length = loader.seq_length
use_rela = getattr(opt, 'use_rela', 0)
try:
tb_summary_writer = tf and tf.summary.FileWriter(opt.checkpoint_path)
except:
print('Set tensorboard error!')
pdb.set_trace()
infos = {}
histories = {}
if opt.start_from is not None:
# open old infos and check if models are compatible
try:
with open(os.path.join(opt.checkpoint_path, 'infos.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.checkpoint_path, 'histories.pkl')):
with open(os.path.join(opt.checkpoint_path,
'histories.pkl')) as f:
histories = cPickle.load(f)
except:
print("Can not load infos.pkl")
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)
# opt.caption_model = 'up_gtssg_sep_self_att_sep'
opt.caption_model = opt.caption_model_to_replace
model = models.setup(opt).cuda()
print('### Model summary below###\n {}\n'.format(str(model)))
model_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('model parameter:{}'.format(model_params))
model.eval()
train_loss = 0
update_lr_flag = True
fake_A_pool_obj = utils.ImagePool(opt.pool_size)
fake_A_pool_rel = utils.ImagePool(opt.pool_size)
fake_A_pool_atr = utils.ImagePool(opt.pool_size)
fake_B_pool_obj = utils.ImagePool(opt.pool_size)
fake_B_pool_rel = utils.ImagePool(opt.pool_size)
fake_B_pool_atr = utils.ImagePool(opt.pool_size)
netD_A_obj = GAN_init_D(opt, Discriminator(opt),
type='netD_A_obj').cuda().train()
netD_A_rel = GAN_init_D(opt, Discriminator(opt),
type='netD_A_rel').cuda().train()
netD_A_atr = GAN_init_D(opt, Discriminator(opt),
type='netD_A_atr').cuda().train()
netD_B_obj = GAN_init_D(opt, Discriminator(opt),
type='netD_B_obj').cuda().train()
netD_B_rel = GAN_init_D(opt, Discriminator(opt),
type='netD_B_rel').cuda().train()
netD_B_atr = GAN_init_D(opt, Discriminator(opt),
type='netD_B_atr').cuda().train()
netG_A_obj = GAN_init_G(opt, Generator(opt),
type='netG_A_obj').cuda().train()
netG_A_rel = GAN_init_G(opt, Generator(opt),
type='netG_A_rel').cuda().train()
netG_A_atr = GAN_init_G(opt, Generator(opt),
type='netG_A_atr').cuda().train()
netG_B_obj = GAN_init_G(opt, Generator(opt),
type='netG_B_obj').cuda().train()
netG_B_rel = GAN_init_G(opt, Generator(opt),
type='netG_B_rel').cuda().train()
netG_B_atr = GAN_init_G(opt, Generator(opt),
type='netG_B_atr').cuda().train()
optimizer_G = utils.build_optimizer(
itertools.chain(netG_A_obj.parameters(), netG_B_obj.parameters(),
netG_A_rel.parameters(), netG_B_rel.parameters(),
netG_A_atr.parameters(), netG_B_atr.parameters()), opt)
optimizer_D = utils.build_optimizer(
itertools.chain(netD_A_obj.parameters(), netD_B_obj.parameters(),
netD_A_rel.parameters(), netD_B_rel.parameters(),
netD_A_atr.parameters(), netD_B_atr.parameters()), opt)
criterionGAN = GANLoss(opt.gan_mode).cuda() # define GAN loss.
criterionCycle = torch.nn.L1Loss()
criterionIdt = torch.nn.L1Loss()
optimizers = []
optimizers.append(optimizer_G)
optimizers.append(optimizer_D)
schedulers = [get_scheduler(opt, optimizer) for optimizer in optimizers]
current_lr = optimizers[0].param_groups[0]['lr']
train_num = 0
update_lr_flag = True
while True:
if update_lr_flag and opt.current_lr >= 1e-4:
# 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
if opt.current_lr >= 1e-4:
utils.set_lr(optimizer, opt.current_lr)
else:
utils.set_lr(optimizer, 1e-4)
update_lr_flag = False
"""
Show the percentage of data loader
"""
if train_num > loader.max_index:
train_num = 0
train_num = train_num + 1
train_precentage = float(train_num) * 100 / float(loader.max_index)
"""
Start training
"""
start = time.time()
# Load data from train split (0)
data = loader.get_batch(opt.train_split)
# print('Read data:', time.time() - start)
torch.cuda.synchronize()
start = time.time()
tmp = [
data['isg_feats'][:, 0, :], data['isg_feats'][:, 1, :],
data['isg_feats'][:, 2, :], data['ssg_feats'][:, 0, :],
data['ssg_feats'][:, 1, :], data['ssg_feats'][:, 2, :]
]
tmp = [
_ if _ is None else torch.from_numpy(_).float().cuda() for _ in tmp
]
real_A_obj, real_A_rel, real_A_atr, real_B_obj, real_B_rel, real_B_atr = tmp
iteration += 1
fake_B_rel = netG_A_rel(real_A_rel)
rec_A_rel = netG_B_rel(fake_B_rel)
idt_B_rel = netG_B_rel(real_A_rel)
fake_A_rel = netG_B_rel(real_B_rel)
rec_B_rel = netG_A_rel(fake_A_rel)
idt_A_rel = netG_A_rel(real_B_rel)
# Obj
fake_B_obj = netG_A_obj(real_A_obj)
rec_A_obj = netG_B_obj(fake_B_obj)
idt_B_obj = netG_B_obj(real_A_obj)
fake_A_obj = netG_B_obj(real_B_obj)
rec_B_obj = netG_A_obj(fake_A_obj)
idt_A_obj = netG_A_obj(real_B_obj)
# Atr
fake_B_atr = netG_A_atr(real_A_atr)
rec_A_atr = netG_B_atr(fake_B_atr)
idt_B_atr = netG_B_atr(real_A_atr)
fake_A_atr = netG_B_atr(real_B_atr)
rec_B_atr = netG_A_atr(fake_A_atr)
idt_A_atr = netG_A_atr(real_B_atr)
domain_A = [
real_A_obj, real_A_rel, real_A_atr, fake_A_obj, fake_A_rel,
fake_A_atr, rec_A_obj, rec_A_rel, rec_A_atr, idt_A_obj, idt_A_rel,
idt_A_atr
]
domain_B = [
real_B_obj, real_B_rel, real_B_atr, fake_B_obj, fake_B_rel,
fake_B_atr, rec_B_obj, rec_B_rel, rec_B_atr, idt_B_obj, idt_B_rel,
idt_B_atr
]
# G_A and G_B
utils.set_requires_grad([
netD_A_obj, netD_A_rel, netD_A_atr, netD_B_obj, netD_B_rel,
netD_B_atr
], False) # Ds require no gradients when optimizing Gs
optimizer_G.zero_grad() # set G_A and G_B's gradients to zero
loss_G = cycle_GAN_backward_G(opt, criterionGAN, criterionCycle,
criterionIdt, netG_A_obj, netG_A_rel,
netG_A_atr, netG_B_obj, netG_B_rel,
netG_B_atr, netD_A_obj, netD_A_rel,
netD_A_atr, netD_B_obj, netD_B_rel,
netD_B_atr, domain_A, domain_B)
loss_G.backward()
optimizer_G.step()
# D_A and D_B
utils.set_requires_grad([
netD_A_obj, netD_A_rel, netD_A_atr, netD_B_obj, netD_B_rel,
netD_B_atr
], True)
optimizer_D.zero_grad() # set D_A and D_B's gradients to zero
loss_D_A = cycle_GAN_backward_D(opt, fake_B_pool_obj, fake_B_pool_rel,
fake_B_pool_atr, netD_A_obj,
netD_A_rel, netD_A_atr, criterionGAN,
real_B_obj, real_B_rel, real_B_atr,
fake_B_obj, fake_B_rel, fake_B_atr)
loss_D_A.backward()
loss_D_B = cycle_GAN_backward_D(opt, fake_A_pool_obj, fake_A_pool_rel,
fake_A_pool_atr, netD_B_obj,
netD_B_rel, netD_B_atr, criterionGAN,
real_A_obj, real_A_rel, real_A_atr,
fake_A_obj, fake_A_rel, fake_A_atr)
loss_D_B.backward()
optimizer_D.step() # update D_A and D_B's weights
end = time.time()
train_loss_G = loss_G.item()
train_loss_D_A = loss_D_A.item()
train_loss_D_B = loss_D_B.item()
print(
"{}/{:.1f}/{}/{}|train_loss={:.3f}|train_loss_G={:.3f}|train_loss_D_A={:.3f}|train_loss_D_B={:.3f}|time/batch = {:.3f}"
.format(opt.id, train_precentage, iteration, epoch, train_loss,
train_loss_G, train_loss_D_A, train_loss_D_B, end - start))
torch.cuda.synchronize()
# Write the training loss summary
if (iteration % opt.losses_log_every == 0) and (iteration != 0):
add_summary_value(tb_summary_writer, 'learning_rate',
opt.current_lr, iteration)
add_summary_value(tb_summary_writer, 'train_loss_G', train_loss_G,
iteration)
add_summary_value(tb_summary_writer, 'train_loss_D_A',
train_loss_D_A, iteration)
add_summary_value(tb_summary_writer, 'train_loss_D_B',
train_loss_D_B, iteration)
# add hype parameters
add_summary_value(tb_summary_writer, 'beam_size', opt.beam_size,
iteration)
add_summary_value(tb_summary_writer, 'lambdaA', opt.lambda_A,
iteration)
add_summary_value(tb_summary_writer, 'lambdaB', opt.lambda_B,
iteration)
add_summary_value(tb_summary_writer, 'pool_size', opt.pool_size,
iteration)
add_summary_value(tb_summary_writer, 'gan_type', opt.gan_type,
iteration)
add_summary_value(tb_summary_writer, 'gan_d_type', opt.gan_d_type,
iteration)
add_summary_value(tb_summary_writer, 'gan_g_type', opt.gan_g_type,
iteration)
if (iteration % opt.save_checkpoint_every == 0) and (iteration != 0):
val_loss = eval_utils_gan.eval_split_gan(opt, model, netG_A_obj,
netG_A_rel, netG_A_atr,
loader, loader_i2t)
val_loss = val_loss.item()
add_summary_value(tb_summary_writer, 'validation loss', val_loss,
iteration)
current_score = -val_loss
best_flag = False
save_id = iteration / opt.save_checkpoint_every
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_D.pth')
torch.save(
{
'epoch': epoch,
'netD_A_atr': netD_A_atr.state_dict(),
'netD_A_obj': netD_A_obj.state_dict(),
'netD_A_rel': netD_A_rel.state_dict(),
'netD_B_atr': netD_B_atr.state_dict(),
'netD_B_obj': netD_B_obj.state_dict(),
'netD_B_rel': netD_B_rel.state_dict()
}, checkpoint_path)
checkpoint_path = os.path.join(opt.checkpoint_path, 'model_G.pth')
torch.save(
{
'epoch': epoch,
'netG_A_atr': netG_A_atr.state_dict(),
'netG_A_obj': netG_A_obj.state_dict(),
'netG_A_rel': netG_A_rel.state_dict(),
'netG_B_atr': netG_B_atr.state_dict(),
'netG_B_obj': netG_B_obj.state_dict(),
'netG_B_rel': netG_B_rel.state_dict()
}, checkpoint_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.pkl'),
'wb') as f:
cPickle.dump(infos, f)
with open(os.path.join(opt.checkpoint_path, 'histories.pkl'),
'wb') as f:
cPickle.dump(histories, f)
if best_flag:
checkpoint_path = os.path.join(opt.checkpoint_path,
'model_D-best.pth')
torch.save(
{
'epoch': epoch,
'netD_A_atr': netD_A_atr.state_dict(),
'netD_A_obj': netD_A_obj.state_dict(),
'netD_A_rel': netD_A_rel.state_dict(),
'netD_B_atr': netD_B_atr.state_dict(),
'netD_B_obj': netD_B_obj.state_dict(),
'netD_B_rel': netD_B_rel.state_dict()
}, checkpoint_path)
checkpoint_path = os.path.join(opt.checkpoint_path,
'model_G-best.pth')
torch.save(
{
'epoch': epoch,
'netG_A_atr': netG_A_atr.state_dict(),
'netG_A_obj': netG_A_obj.state_dict(),
'netG_A_rel': netG_A_rel.state_dict(),
'netG_B_atr': netG_B_atr.state_dict(),
'netG_B_obj': netG_B_obj.state_dict(),
'netG_B_rel': netG_B_rel.state_dict()
}, checkpoint_path)
print("model saved to {}".format(checkpoint_path))
with open(os.path.join(opt.checkpoint_path, 'infos-best.pkl'),
'wb') as f:
cPickle.dump(infos, f)
# Update the iteration and epoch
if data['bounds']['wrapped']:
# current_lr = update_learning_rate(schedulers, optimizers)
epoch += 1
update_lr_flag = True
# make evaluation on validation set, and save model
# lang_stats_isg = eval_utils_gan.eval_split_i2t(opt, model, netG_A_obj, netG_A_rel, netG_A_atr, loader, loader_i2t)
lang_stats_isg = eval_utils_gan.eval_split_g2t(
opt, model, netG_A_obj, netG_A_rel, netG_A_atr, loader,
loader_i2t)
if lang_stats_isg is not None:
for k, v in lang_stats_isg.items():
add_summary_value(tb_summary_writer, k, v, iteration)
# Stop if reaching max epochs
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
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
model.train()
crit = utils.LanguageModelCriterion()
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:
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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
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()
loss = crit(model(fc_feats, att_feats, labels), labels[:, 1:],
masks[:, 1:])
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, 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)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
lr_history[iteration] = opt.current_lr
ss_prob_history[iteration] = model.ss_prob
# Stop if reaching max epochs
if epoch >= 8:
break
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()
back_model = models.setup(opt, reverse=True)
back_model.cuda()
update_lr_flag = True
# Assure in training mode
model.train()
back_model.train()
crit = utils.LanguageModelCriterion()
all_param = chain(model.parameters(), back_model.parameters())
optimizer = optim.Adam(all_param,
lr=opt.learning_rate,
weight_decay=opt.weight_decay)
# Load the optimizer
if vars(opt).get('start_from', None) is not None:
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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
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()
reverse_labels = np.flip(data['labels'], 1).copy()
reverse_masks = np.flip(data['masks'], 1).copy()
tmp = [
data['fc_feats'], data['att_feats'], data['labels'],
reverse_labels, data['masks'], reverse_masks
]
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
fc_feats, att_feats, labels, reverse_labels, masks, reverse_masks = tmp
optimizer.zero_grad()
out, states = model(fc_feats, att_feats, labels)
back_out, back_states = back_model(fc_feats, att_feats, reverse_labels)
idx = [i for i in range(back_states.size()[1] - 1, -1, -1)]
idx = torch.LongTensor(idx)
idx = Variable(idx).cuda()
invert_backstates = back_states.index_select(1, idx)
loss = crit(out, labels[:, 1:], masks[:, 1:])
back_loss = crit(back_out, reverse_labels[:, :-1],
reverse_masks[:, :-1])
invert_backstates = invert_backstates.detach()
l2_loss = ((states - invert_backstates)**2).mean()
all_loss = loss + 1.5 * l2_loss + back_loss
all_loss.backward()
#back_loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_l2_loss = l2_loss.data[0]
train_loss = loss.data[0]
train_all_loss = all_loss.data[0]
train_back_loss = back_loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, l2_loss = {:.3f}, back_loss = {:.3f}, all_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, train_l2_loss, train_back_loss, train_all_loss, 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, 'l2_loss', train_l2_loss,
iteration)
add_summary_value(tf_summary_writer, 'all_loss',
train_all_loss, iteration)
add_summary_value(tf_summary_writer, 'back_loss',
train_back_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)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
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)
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(opt):
opt.use_att = utils.if_use_att(opt.caption_model)
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.maxlen_sen
opt.inc_seg = loader.inc_seg
opt.seg_ix = loader.seg_ix
tf_summary_writer = tf and tf.summary.FileWriter(opt.checkpoint_path)
infos = {}
histories = {}
score_list = []
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)
best_val_score = None
best_val_score = {}
score_splits = ['val', 'test']
score_type = ['Bleu_4', 'METEOR', 'CIDEr']
for split_i in score_splits:
for score_item in score_type:
if split_i not in best_val_score.keys():
best_val_score[split_i] = {}
best_val_score[split_i][score_item] = 0.0
if opt.load_best_score == 1:
best_val_score = infos.get('best_val_score', best_val_score)
model = models.setup(opt)
device_ids = [0, 1]
torch.cuda.set_device(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)
model = model.cuda()
update_lr_flag = True
# Assure in training mode
model.module.train()
crit = utils.LanguageModelCriterion()
optimizer = optim.Adam(model.module.parameters(),
lr=opt.learning_rate,
weight_decay=opt.weight_decay)
#optimizer = nn.DataParallel(optimizer, device_ids=device_ids)
# Load the optimizer
if vars(opt).get('start_from', None) is not None:
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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.module.ss_prob = opt.ss_prob
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['labels'], data['x_phrase_mask_0'], data['x_phrase_mask_1'], \
data['label_masks'], data['salicy_seg'], data['seg_mask']]
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
fc_feats, seq, phrase_mask_0, phrase_mask_1, masks, salicy_seg, seg_mask = tmp
optimizer.zero_grad()
remove_len = 2
outputs, alphas = model.module(fc_feats, seq, phrase_mask_0,
phrase_mask_1, masks, seg_mask,
remove_len)
loss = crit(outputs, seq[remove_len:, :].permute(1, 0),
masks[remove_len:, :].permute(1, 0))
alphas = alphas.permute(1, 0, 2)
salicy_seg = salicy_seg[:, :, :]
seg_mask = seg_mask[:, :]
if opt.salicy_hard == False:
if opt.salicy_loss_type == 'l2':
salicy_loss = (((((salicy_seg * seg_mask[:, :, None] -
alphas * seg_mask[:, :, None])**2).sum(0)
).sum(-1))**(0.5)).mean()
if opt.salicy_loss_type == 'kl':
#alphas: len_sen, batch_size, num_frame
salicy_loss = kullback_leibler2(
alphas * seg_mask[:, :, None],
salicy_seg * seg_mask[:, :, None])
salicy_loss = (((salicy_loss *
seg_mask[:, :, None]).sum(-1)).sum(0)).mean()
elif opt.salicy_hard == True:
#salicy len_sen, batch_size, num_frame
salicy_loss = -torch.log((alphas * salicy_seg).sum(-1) + 1e-8)
#salicy_loss len_sen, batch_size
salicy_loss = ((salicy_loss * seg_mask).sum(0)).mean()
loss = loss + opt.salicy_alpha * salicy_loss
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, 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.module.ss_prob, iteration)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
lr_history[iteration] = opt.current_lr
ss_prob_history[iteration] = model.module.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.dataset,
'remove_len': remove_len
}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats, score_list_i = eval_utils.eval_split(
model.module, crit, loader, eval_kwargs)
score_list.append(score_list_i)
np.savetxt('./save/train_valid_test.txt', score_list, fmt='%.3f')
# Write validation result into summary
if tf is not None:
add_summary_value(tf_summary_writer, 'validation loss',
val_loss, iteration)
for k in lang_stats.keys():
for v in lang_stats[k].keys():
add_summary_value(tf_summary_writer, k + v,
lang_stats[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['val']['CIDEr']
else:
current_score = -val_loss
best_flag = {}
for split_i in score_splits:
for score_item in score_type:
if split_i not in best_flag.keys():
best_flag[split_i] = {}
best_flag[split_i][score_item] = False
if True: # if true
for split_i in score_splits:
for score_item in score_type:
if best_val_score is None or lang_stats[split_i][
score_item] > best_val_score[split_i][
score_item]:
best_val_score[split_i][score_item] = lang_stats[
split_i][score_item]
best_flag[split_i][score_item] = True
checkpoint_path = os.path.join(opt.checkpoint_path,
'model.pth')
torch.save(model.module.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)
for split_i in score_splits:
for score_item in score_type:
if best_flag[split_i][score_item]:
checkpoint_path = os.path.join(
opt.checkpoint_path, 'model-best_' + split_i +
'_' + score_item + '.pth')
torch.save(model.module.state_dict(),
checkpoint_path)
print("model saved to {}".format(checkpoint_path))
with open(
os.path.join(
opt.checkpoint_path,
'infos_' + split_i + '_' + score_item +
'_' + 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(opt):
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
opt.use_fc = utils.if_use_fc(opt.caption_model)
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
infos = load_info(opt)
iteration = infos.get('iter', 0)
epoch = infos.get('epoch', 0)
val_result_history = infos.get('val_result_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)
# Define and load model, optimizer, critics
decoder = setup(opt).train().cuda()
if opt.label_smoothing > 0:
crit = utils.LabelSmoothing(smoothing=opt.label_smoothing).cuda()
else:
crit = utils.LanguageModelCriterion().cuda()
# crit = utils.LanguageModelCriterion().cuda()
rl_crit = utils.RewardCriterion().cuda()
if opt.reduce_on_plateau:
optimizer = utils.build_optimizer(decoder.parameters(), opt)
optimizer = utils.ReduceLROnPlateau(optimizer, factor=0.5, patience=3)
else:
optimizer = utils.build_optimizer(decoder.parameters(), opt)
# optimizer = utils.build_optimizer(decoder.parameters(), opt)
models = {'decoder': decoder}
optimizers = {'decoder': optimizer}
save_nets_structure(models, opt)
load_checkpoint(models, optimizers, opt)
print('opt', opt)
epoch_done = True
sc_flag = False
while True:
if epoch_done:
# 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)
decoder.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_scorer(opt.cached_tokens)
else:
sc_flag = False
epoch_done = False
# 1. fetch a batch of data from train split
data = loader.get_batch('train')
tmp = [
data['fc_feats'], data['att_feats'], data['labels'], data['tags'],
data['masks'], data['att_masks'], data['verbs']
]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
fc_feats, att_feats, labels, tags, masks, att_masks, weak_relas = tmp
vrg_data = {key: data['vrg_data'][key] if data['vrg_data'][key] is None \
else torch.from_numpy(data['vrg_data'][key]).cuda() for key in data['vrg_data']}
# 2. Forward model and compute loss
torch.cuda.synchronize()
optimizer.zero_grad()
if not sc_flag:
out = decoder(vrg_data, fc_feats, att_feats, labels, weak_relas,
att_masks)
loss_words = crit(out[0], labels[:, 1:], masks[:, 1:])
loss_tags = crit(out[1], tags[:, 1:], masks[:, 1:])
loss = loss_words + loss_tags * 0.15
else:
gen_result, sample_logprobs, core_args = decoder(
vrg_data,
fc_feats,
att_feats,
weak_relas,
att_masks,
opt={
'sample_max': 0,
'return_core_args': True
},
mode='sample')
reward = get_self_critical_reward(decoder, core_args, vrg_data,
fc_feats, att_feats, weak_relas,
att_masks, data, gen_result, opt)
loss = rl_crit(sample_logprobs, gen_result.data,
torch.from_numpy(reward).float().cuda())
# 3. Update model
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.item()
torch.cuda.synchronize()
# Update the iteration and epoch
iteration += 1
# Write the training loss summary
if (iteration % opt.log_loss_every == 0):
# logging log
logger.info("{} ({}), loss: {:.3f}".format(iteration, epoch,
train_loss))
tb.add_values('loss', {'train': train_loss}, iteration)
if data['bounds']['wrapped']:
epoch += 1
epoch_done = True
# Make evaluation and save checkpoint
if (opt.save_checkpoint_every > 0
and iteration % opt.save_checkpoint_every
== 0) or (opt.save_checkpoint_every == -1 and epoch_done):
# eval model
eval_kwargs = {
'split': 'val',
'dataset': opt.input_json,
'expand_features': False
}
eval_kwargs.update(vars(opt))
predictions, lang_stats = eval_utils.eval_split(
decoder, loader, eval_kwargs)
if opt.reduce_on_plateau:
assert 'CIDEr' in lang_stats, 'Error: cider should be in eval list'
optimizer.scheduler_step(-lang_stats['CIDEr'])
# log val results
if not lang_stats is None:
logger.info("Scores: {}".format(lang_stats))
tb.add_values('scores', lang_stats, epoch)
val_result_history[epoch] = {
'lang_stats': lang_stats,
'predictions': predictions
}
# Save model if is improving on validation result
current_score = 0 if lang_stats is None else lang_stats['CIDEr']
best_flag = False
if best_val_score is None or current_score > best_val_score:
best_val_score = current_score
best_flag = True
# 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()
infos['val_result_history'] = val_result_history
save_checkpoint(models, optimizers, infos, best_flag, opt)
# Stop if reaching max epochs
if epoch > opt.max_epochs and opt.max_epochs != -1:
break
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_size1", "rnn_size2",
"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()
# model.set_mode('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:
model.train()
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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+val')
# print('Read data:', time.time() - start)
torch.cuda.synchronize()
start = time.time()
tmp = [
data['fc_feats'], data['att_feats'], data['num_bbox'],
data['labels'], data['masks']
]
tmp = [
Variable(torch.from_numpy(_).float(), requires_grad=False).cuda()
for _ in tmp
]
fc_feats, att_feats, num_bbox, labels, masks = tmp
labels = labels.long()
optimizer.zero_grad()
if not sc_flag:
loss = crit(model(fc_feats, att_feats, num_bbox, labels),
labels[:, 1:], masks[:, 1:])
# loss = crit(model(fc_feats, att_feats, labels), labels[:,1:], masks[:,1:])
else:
gen_result, sample_logprobs = model.sample(fc_feats, att_feats,
num_bbox,
{'sample_max': 0})
reward = get_self_critical_reward(model, fc_feats, att_feats,
num_bbox, 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:
if (iteration % 100 == 0):
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f} lr={}" \
.format(iteration, epoch, train_loss, end - start, opt.current_lr ))
else:
if (iteration % 100 == 0):
print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f} lr={}" \
.format(iteration, epoch, np.mean(reward[:,0]), end - start, opt.current_lr ))
# 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,
'val_ref_path': opt.val_ref_path,
'raw_val_anno_path': opt.raw_val_anno_path
}
eval_kwargs.update(vars(opt))
# predictions, lang_stats = eval_utils.eval_split(model, crit, loader, eval_kwargs)
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(opt):
opt['source'].use_att = utils.if_use_att(opt['source'].caption_model)
loader = {}
loader['source'] = DataLoader(opt['source'])
opt['source'].vocab_size = loader[
'source'].vocab_size # vocab and seq_length all follows the source dataset
opt['source'].seq_length = loader['source'].seq_length
loader['target'] = DataLoader(opt['target'],
target=True,
opt_extra=opt['source'])
tf_summary_writer = tf and tf.summary.FileWriter(
opt['source'].checkpoint_path)
infos = {}
histories = {}
if opt['source'].start_from is not None: # the state of previous training session
# open old infos and check if models are compatible
with open(
os.path.join(opt['source'].start_from,
'infos_' + opt['source'].id + '.pkl')) as f:
infos = cPickle.load(f)
saved_model_opt = infos['opt']['source']
need_be_same = [
"caption_model", "rnn_type", "rnn_size", "num_layers"
]
for checkme in need_be_same: # the key arguments of RNN model should keep the same
assert vars(saved_model_opt)[checkme] == vars(
opt['source']
)[checkme], "Command line argument and saved model disagree on '%s' " % checkme
if os.path.isfile(
os.path.join(opt['source'].start_from,
'histories_' + opt['source'].id + '.pkl')):
with open(
os.path.join(opt['source'].start_from, 'histories_' +
opt['source'].id + '.pkl')) as f:
histories = cPickle.load(
f
) # it is the histories file that be used in the following training
iteration = infos.get(
'iter',
0) # obtain the iteration number, if not defined, then return 0
epoch = infos.get(
'epoch', 0) # obtain the epoch number, if not defined, then return 0
val_result_history = histories.get('val_result_history',
{}) # if not defined, return {}
loss_history = histories.get('loss_history', {})
lr_history = histories.get('lr_history', {})
ss_prob_history = histories.get('ss_prob_history', {})
loader['source'].iterator = infos.get(
'iterator',
loader['source'].iterator) # if not defined, return loader.iterator
loader['source'].split_ix = infos.get('split_ix',
loader['source'].split_ix)
if opt['source'].load_best_score == 1:
best_val_score = infos.get('best_val_score', None)
model = models.setup(
opt['source']
) # opt contains the model to use, here is the caption training model, the definitions are in __init__.py
model.cuda()
update_lr_flag = True
# Assure in training mode
model.train()
crit = utils.LanguageModelCriterion() # for later learning
optimizer = optim.Adam(model.parameters(),
lr=opt['source'].learning_rate,
weight_decay=opt['source'].weight_decay)
# Load the optimizer, if training from another status instead of scratch
if vars(opt['source']).get('start_from', None) is not None:
optimizer.load_state_dict(
torch.load(os.path.join(opt['source'].start_from,
'optimizer.pth')))
while True:
if update_lr_flag:
# Assign the learning rate
if epoch > opt['source'].learning_rate_decay_start and opt[
'source'].learning_rate_decay_start >= 0:
frac = (epoch - opt['source'].learning_rate_decay_start
) // opt['source'].learning_rate_decay_every
decay_factor = opt['source'].learning_rate_decay_rate**frac
opt['source'].current_lr = opt[
'source'].learning_rate * decay_factor
utils.set_lr(optimizer,
opt['source'].current_lr) # set the decayed rate
else:
opt['source'].current_lr = opt['source'].learning_rate
# Assign the scheduled sampling prob
if epoch > opt['source'].scheduled_sampling_start and opt[
'source'].scheduled_sampling_start >= 0:
frac = (epoch - opt['source'].scheduled_sampling_start
) // opt['source'].scheduled_sampling_increase_every
opt['source'].ss_prob = min(
opt['source'].scheduled_sampling_increase_prob * frac,
opt['source'].scheduled_sampling_max_prob)
model.ss_prob = opt['source'].ss_prob
update_lr_flag = False
start = time.time()
# Load data from train split (0)
data = loader['source'].get_batch()
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()
loss = crit(model(fc_feats, att_feats, labels), labels[:, 1:],
masks[:, 1:])
loss.backward()
utils.clip_gradient(optimizer, opt['source'].grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, 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['source'].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['source'].current_lr, iteration)
add_summary_value(tf_summary_writer, 'scheduled_sampling_prob',
model.ss_prob, iteration)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
lr_history[iteration] = opt['source'].current_lr
ss_prob_history[iteration] = model.ss_prob
# make evaluation on validation set, and save model
if (iteration % opt['source'].save_checkpoint_every == 0):
# eval model
eval_kwargs = {'dataset': opt['target'].input_target}
eval_kwargs.update(
vars(opt['target'])
) # the final version of eval_kwargs is a dict same as opt added with the previous key
val_loss, predictions, lang_stats = eval_utils.eval_split(
model, crit, loader['target'], eval_kwargs)
# Write validation result into summary
if tf is not None:
add_summary_value(tf_summary_writer, 'validation loss',
val_loss, iteration)
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['target'].language_eval == 1:
current_score = lang_stats['CIDEr'] # only based on 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: # choose the best model based on the validation score
best_val_score = current_score
best_flag = True
checkpoint_path = os.path.join(
opt['source'].checkpoint_path,
'model' + str(iteration) + '.pth')
torch.save(model.state_dict(), checkpoint_path)
print("model saved to {}".format(checkpoint_path))
optimizer_path = os.path.join(
opt['source'].checkpoint_path,
'optimizer' + str(iteration) + '.pth')
torch.save(optimizer.state_dict(), optimizer_path)
# Dump miscalleous informations
infos['iter'] = iteration
infos['epoch'] = epoch
infos['iterator'] = loader['source'].iterator
infos['split_ix'] = loader['source'].split_ix
infos['best_val_score'] = best_val_score
infos['opt'] = opt
infos['vocab'] = loader['source'].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['source'].checkpoint_path, 'infos_' +
opt['source'].id + str(iteration) + '.pkl'),
'wb') as f:
cPickle.dump(infos, f)
with open(
os.path.join(
opt['source'].checkpoint_path, 'histories_' +
opt['source'].id + str(iteration) + '.pkl'),
'wb') as f:
cPickle.dump(histories, f)
if best_flag:
checkpoint_path = os.path.join(
opt['source'].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['source'].checkpoint_path,
'infos_' + opt['source'].id + '-best.pkl'),
'wb') as f:
cPickle.dump(infos, f)
# Stop if reaching max epochs
if epoch >= opt['source'].max_epochs and opt['source'].max_epochs != -1:
break
def train(opt):
exclude_opt = [
'training_mode', 'tap_epochs', 'cg_epochs', 'tapcg_epochs', 'lr',
'learning_rate_decay_start', 'learning_rate_decay_every',
'learning_rate_decay_rate', 'self_critical_after',
'save_checkpoint_every', 'id', "pretrain", "pretrain_path", "debug",
"save_all_checkpoint", "min_epoch_when_save"
]
save_folder, logger, tf_writer = build_floder_and_create_logger(opt)
saved_info = {'best': {}, 'last': {}, 'history': {}}
is_continue = opt.start_from != None
if is_continue:
infos_path = os.path.join(save_folder, 'info.pkl')
with open(infos_path) as f:
logger.info('load info from {}'.format(infos_path))
saved_info = cPickle.load(f)
pre_opt = saved_info[opt.start_from_mode]['opt']
if vars(opt).get("no_exclude_opt", False):
exclude_opt = []
for opt_name in vars(pre_opt).keys():
if (not opt_name in exclude_opt):
vars(opt).update({opt_name: vars(pre_opt).get(opt_name)})
if vars(pre_opt).get(opt_name) != vars(opt).get(opt_name):
print('change opt: {} from {} to {}'.format(
opt_name,
vars(pre_opt).get(opt_name),
vars(opt).get(opt_name)))
opt.use_att = utils.if_use_att(opt.caption_model)
loader = DataLoader(opt)
opt.CG_vocab_size = loader.vocab_size
opt.CG_seq_length = loader.seq_length
# init training option
epoch = saved_info[opt.start_from_mode].get('epoch', 0)
iteration = saved_info[opt.start_from_mode].get('iter', 0)
best_val_score = saved_info[opt.start_from_mode].get('best_val_score', 0)
val_result_history = saved_info['history'].get('val_result_history', {})
loss_history = saved_info['history'].get('loss_history', {})
lr_history = saved_info['history'].get('lr_history', {})
loader.iterators = saved_info[opt.start_from_mode].get(
'iterators', loader.iterators)
loader.split_ix = saved_info[opt.start_from_mode].get(
'split_ix', loader.split_ix)
opt.current_lr = vars(opt).get('current_lr', opt.lr)
opt.m_batch = vars(opt).get('m_batch', 1)
# create a tap_model,fusion_model,cg_model
tap_model = models.setup_tap(opt)
lm_model = CaptionGenerator(opt)
cg_model = lm_model
if is_continue:
if opt.start_from_mode == 'best':
model_pth = torch.load(os.path.join(save_folder, 'model-best.pth'))
elif opt.start_from_mode == 'last':
model_pth = torch.load(
os.path.join(save_folder,
'model_iter_{}.pth'.format(iteration)))
assert model_pth['iteration'] == iteration
logger.info('Loading pth from {}, iteration:{}'.format(
save_folder, iteration))
tap_model.load_state_dict(model_pth['tap_model'])
cg_model.load_state_dict(model_pth['cg_model'])
elif opt.pretrain:
print('pretrain {} from {}'.format(opt.pretrain, opt.pretrain_path))
model_pth = torch.load(opt.pretrain_path)
if opt.pretrain == 'tap':
tap_model.load_state_dict(model_pth['tap_model'])
elif opt.pretrain == 'cg':
cg_model.load_state_dict(model_pth['cg_model'])
elif opt.pretrain == 'tap_cg':
tap_model.load_state_dict(model_pth['tap_model'])
cg_model.load_state_dict(model_pth['cg_model'])
else:
assert 1 == 0, 'opt.pretrain error'
tap_model.cuda()
tap_model.train() # Assure in training mode
tap_crit = utils.TAPModelCriterion()
tap_optimizer = optim.Adam(tap_model.parameters(),
lr=opt.lr,
weight_decay=opt.weight_decay)
cg_model.cuda()
cg_model.train()
cg_optimizer = optim.Adam(cg_model.parameters(),
lr=opt.lr,
weight_decay=opt.weight_decay)
cg_crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
cg_optimizer = optim.Adam(cg_model.parameters(),
lr=opt.lr,
weight_decay=opt.weight_decay)
allmodels = [tap_model, cg_model]
optimizers = [tap_optimizer, cg_optimizer]
if is_continue:
tap_optimizer.load_state_dict(model_pth['tap_optimizer'])
cg_optimizer.load_state_dict(model_pth['cg_optimizer'])
update_lr_flag = True
loss_sum = np.zeros(5)
bad_video_num = 0
best_epoch = epoch
start = time.time()
print_opt(opt, allmodels, logger)
logger.info('\nStart training')
# set a var to indicate what to train in current iteration: "tap", "cg" or "tap_cg"
flag_training_whats = get_training_list(opt, logger)
# Iteration begin
while True:
if update_lr_flag:
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.lr * decay_factor
else:
opt.current_lr = opt.lr
for optimizer in optimizers:
utils.set_lr(optimizer, opt.current_lr)
if opt.self_critical_after != -1 and epoch >= opt.self_critical_after:
sc_flag = True
init_scorer(None)
else:
sc_flag = False
update_lr_flag = False
flag_training_what = flag_training_whats[epoch]
if opt.training_mode == "alter2":
flag_training_what = flag_training_whats[iteration]
# get data
data = loader.get_batch('train')
if opt.debug:
print('vid:', data['vid'])
print('info:', data['infos'])
torch.cuda.synchronize()
if (data["proposal_num"] <= 0) or (data['fc_feats'].shape[0] <= 1):
bad_video_num += 1 # print('vid:{} has no good proposal.'.format(data['vid']))
continue
ind_select_list, soi_select_list, cg_select_list, sampled_ids, = data[
'ind_select_list'], data['soi_select_list'], data[
'cg_select_list'], data['sampled_ids']
if flag_training_what == 'cg' or flag_training_what == 'gt_tap_cg':
ind_select_list = data['gts_ind_select_list']
soi_select_list = data['gts_soi_select_list']
cg_select_list = data['gts_cg_select_list']
tmp = [
data['fc_feats'], data['att_feats'], data['lda_feats'],
data['tap_labels'], data['tap_masks_for_loss'],
data['cg_labels'][cg_select_list],
data['cg_masks'][cg_select_list], data['w1']
]
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
c3d_feats, att_feats, lda_feats, tap_labels, tap_masks_for_loss, cg_labels, cg_masks, w1 = tmp
if (iteration - 1) % opt.m_batch == 0:
tap_optimizer.zero_grad()
cg_optimizer.zero_grad()
tap_feats, pred_proposals = tap_model(c3d_feats)
tap_loss = tap_crit(pred_proposals, tap_masks_for_loss, tap_labels, w1)
loss_sum[0] = loss_sum[0] + tap_loss.item()
# Backward Propagation
if flag_training_what == 'tap':
tap_loss.backward()
utils.clip_gradient(tap_optimizer, opt.grad_clip)
if iteration % opt.m_batch == 0:
tap_optimizer.step()
else:
if not sc_flag:
pred_captions = cg_model(tap_feats,
c3d_feats,
lda_feats,
cg_labels,
ind_select_list,
soi_select_list,
mode='train')
cg_loss = cg_crit(pred_captions, cg_labels[:, 1:],
cg_masks[:, 1:])
else:
gen_result, sample_logprobs, greedy_res = cg_model(
tap_feats,
c3d_feats,
lda_feats,
cg_labels,
ind_select_list,
soi_select_list,
mode='train_rl')
sentence_info = data['sentences_batch'] if (
flag_training_what != 'cg'
and flag_training_what != 'gt_tap_cg'
) else data['gts_sentences_batch']
reward = get_self_critical_reward2(
greedy_res, (data['vid'], sentence_info),
gen_result,
vocab=loader.get_vocab(),
opt=opt)
cg_loss = rl_crit(sample_logprobs, gen_result,
torch.from_numpy(reward).float().cuda())
loss_sum[1] = loss_sum[1] + cg_loss.item()
if flag_training_what == 'cg' or flag_training_what == 'gt_tap_cg' or flag_training_what == 'LP_cg':
cg_loss.backward()
utils.clip_gradient(cg_optimizer, opt.grad_clip)
if iteration % opt.m_batch == 0:
cg_optimizer.step()
if flag_training_what == 'gt_tap_cg':
utils.clip_gradient(tap_optimizer, opt.grad_clip)
if iteration % opt.m_batch == 0:
tap_optimizer.step()
elif flag_training_what == 'tap_cg':
total_loss = opt.lambda1 * tap_loss + opt.lambda2 * cg_loss
total_loss.backward()
utils.clip_gradient(tap_optimizer, opt.grad_clip)
utils.clip_gradient(cg_optimizer, opt.grad_clip)
if iteration % opt.m_batch == 0:
tap_optimizer.step()
cg_optimizer.step()
loss_sum[2] = loss_sum[2] + total_loss.item()
torch.cuda.synchronize()
# Updating epoch num
iteration += 1
if data['bounds']['wrapped']:
epoch += 1
update_lr_flag = True
# Print losses, Add to summary
if iteration % opt.losses_log_every == 0:
end = time.time()
losses = np.round(loss_sum / opt.losses_log_every, 3)
logger.info(
"iter {} (epoch {}, lr {}), avg_iter_loss({}) = {}, time/batch = {:.3f}, bad_vid = {:.3f}" \
.format(iteration, epoch, opt.current_lr, flag_training_what, losses,
(end - start) / opt.losses_log_every,
bad_video_num))
tf_writer.add_scalar('lr', opt.current_lr, iteration)
tf_writer.add_scalar('train_tap_loss', losses[0], iteration)
tf_writer.add_scalar('train_tap_prop_loss', losses[3], iteration)
tf_writer.add_scalar('train_tap_bound_loss', losses[4], iteration)
tf_writer.add_scalar('train_cg_loss', losses[1], iteration)
tf_writer.add_scalar('train_total_loss', losses[2], iteration)
if sc_flag and (not flag_training_what == 'tap'):
tf_writer.add_scalar('avg_reward', np.mean(reward[:, 0]),
iteration)
loss_history[iteration] = losses
lr_history[iteration] = opt.current_lr
loss_sum = np.zeros(5)
start = time.time()
bad_video_num = 0
# Evaluation, and save model
if (iteration % opt.save_checkpoint_every
== 0) and (epoch >= opt.min_epoch_when_save):
eval_kwargs = {
'split': 'val',
'val_all_metrics': 0,
'topN': 100,
}
eval_kwargs.update(vars(opt))
# eval_kwargs['num_vids_eval'] = int(491)
eval_kwargs['topN'] = 100
eval_kwargs2 = {
'split': 'val',
'val_all_metrics': 1,
'num_vids_eval': 4917,
}
eval_kwargs2.update(vars(opt))
if not opt.num_vids_eval:
eval_kwargs['num_vids_eval'] = int(4917.)
eval_kwargs2['num_vids_eval'] = 4917
crits = [tap_crit, cg_crit]
pred_json_path_T = os.path.join(save_folder, 'pred_sent',
'pred_num{}_iter{}.json')
# if 'alter' in opt.training_mode:
if flag_training_what == 'tap':
eval_kwargs['topN'] = 1000
predictions, eval_score, val_loss = eval_utils.eval_split(
allmodels,
crits,
loader,
pred_json_path_T.format(eval_kwargs['num_vids_eval'],
iteration),
eval_kwargs,
flag_eval_what='tap')
else:
if vars(opt).get('fast_eval_cg', False) == False:
predictions, eval_score, val_loss = eval_utils.eval_split(
allmodels,
crits,
loader,
pred_json_path_T.format(eval_kwargs['num_vids_eval'],
iteration),
eval_kwargs,
flag_eval_what='tap_cg')
predictions2, eval_score2, val_loss2 = eval_utils.eval_split(
allmodels,
crits,
loader,
pred_json_path_T.format(eval_kwargs2['num_vids_eval'],
iteration),
eval_kwargs2,
flag_eval_what='cg')
if (not vars(opt).get('fast_eval_cg', False)
== False) or (not vars(opt).get(
'fast_eval_cg_top10', False) == False):
eval_score = eval_score2
val_loss = val_loss2
predictions = predictions2
# else:
# predictions, eval_score, val_loss = eval_utils.eval_split(allmodels, crits, loader, pred_json_path,
# eval_kwargs,
# flag_eval_what=flag_training_what)
f_f1 = lambda x, y: 2 * x * y / (x + y)
f1 = f_f1(eval_score['Recall'], eval_score['Precision']).mean()
if flag_training_what != 'tap': # if only train tap, use the mean of precision and recall as final score
current_score = np.array(eval_score['METEOR']).mean() * 100
else: # if train tap_cg, use avg_meteor as final score
current_score = f1
for model in allmodels:
for name, param in model.named_parameters():
tf_writer.add_histogram(name,
param.clone().cpu().data.numpy(),
iteration,
bins=10)
if param.grad is not None:
tf_writer.add_histogram(
name + '_grad',
param.grad.clone().cpu().data.numpy(),
iteration,
bins=10)
tf_writer.add_scalar('val_tap_loss', val_loss[0], iteration)
tf_writer.add_scalar('val_cg_loss', val_loss[1], iteration)
tf_writer.add_scalar('val_tap_prop_loss', val_loss[3], iteration)
tf_writer.add_scalar('val_tap_bound_loss', val_loss[4], iteration)
tf_writer.add_scalar('val_total_loss', val_loss[2], iteration)
tf_writer.add_scalar('val_score', current_score, iteration)
if flag_training_what != 'tap':
tf_writer.add_scalar('val_score_gt_METEOR',
np.array(eval_score2['METEOR']).mean(),
iteration)
tf_writer.add_scalar('val_score_gt_Bleu_4',
np.array(eval_score2['Bleu_4']).mean(),
iteration)
tf_writer.add_scalar('val_score_gt_CIDEr',
np.array(eval_score2['CIDEr']).mean(),
iteration)
tf_writer.add_scalar('val_recall', eval_score['Recall'].mean(),
iteration)
tf_writer.add_scalar('val_precision',
eval_score['Precision'].mean(), iteration)
tf_writer.add_scalar('f1', f1, iteration)
val_result_history[iteration] = {
'val_loss': val_loss,
'eval_score': eval_score
}
if flag_training_what == 'tap':
logger.info(
'Validation the result of iter {}, score(f1/meteor):{},\n all:{}'
.format(iteration, current_score, eval_score))
else:
mean_score = {
k: np.array(v).mean()
for k, v in eval_score.items()
}
gt_mean_score = {
k: np.array(v).mean()
for k, v in eval_score2.items()
}
metrics = ['Bleu_4', 'CIDEr', 'METEOR', 'ROUGE_L']
gt_avg_score = np.array([
v for metric, v in gt_mean_score.items()
if metric in metrics
]).sum()
logger.info(
'Validation the result of iter {}, score(f1/meteor):{},\n all:{}\n mean:{} \n\n gt:{} \n mean:{}\n avg_score: {}'
.format(iteration, current_score, eval_score, mean_score,
eval_score2, gt_mean_score, gt_avg_score))
# Save model .pth
saved_pth = {
'iteration': iteration,
'cg_model': cg_model.state_dict(),
'tap_model': tap_model.state_dict(),
'cg_optimizer': cg_optimizer.state_dict(),
'tap_optimizer': tap_optimizer.state_dict(),
}
if opt.save_all_checkpoint:
checkpoint_path = os.path.join(
save_folder, 'model_iter_{}.pth'.format(iteration))
else:
checkpoint_path = os.path.join(save_folder, 'model.pth')
torch.save(saved_pth, checkpoint_path)
logger.info('Save model at iter {} to checkpoint file {}.'.format(
iteration, checkpoint_path))
# save info.pkl
if current_score > best_val_score:
best_val_score = current_score
best_epoch = epoch
saved_info['best'] = {
'opt': opt,
'iter': iteration,
'epoch': epoch,
'iterators': loader.iterators,
'flag_training_what': flag_training_what,
'split_ix': loader.split_ix,
'best_val_score': best_val_score,
'vocab': loader.get_vocab(),
}
best_checkpoint_path = os.path.join(save_folder,
'model-best.pth')
torch.save(saved_pth, best_checkpoint_path)
logger.info(
'Save Best-model at iter {} to checkpoint file.'.format(
iteration))
saved_info['last'] = {
'opt': opt,
'iter': iteration,
'epoch': epoch,
'iterators': loader.iterators,
'flag_training_what': flag_training_what,
'split_ix': loader.split_ix,
'best_val_score': best_val_score,
'vocab': loader.get_vocab(),
}
saved_info['history'] = {
'val_result_history': val_result_history,
'loss_history': loss_history,
'lr_history': lr_history,
}
with open(os.path.join(save_folder, 'info.pkl'), 'w') as f:
cPickle.dump(saved_info, f)
logger.info('Save info to info.pkl')
# Stop criterion
if epoch >= len(flag_training_whats):
tf_writer.close()
break
def train(opt):
# Deal with feature things before anything
opt.use_att = utils.if_use_att(opt.caption_model)
ac = 0
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
tb_summary_writer = tb and tb.SummaryWriter(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.checkpoint_path, 'infos_' + opt.id + format(int(opt.start_from),'04') + '.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.checkpoint_path, 'histories_' + opt.id + format(int(opt.start_from),'04') + '.pkl')):
with open(os.path.join(opt.checkpoint_path, 'histories_' + opt.id + format(int(opt.start_from),'04') + '.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).cuda()
#dp_model = torch.nn.DataParallel(model)
#dp_model = torch.nn.DataParallel(model, [0,2,3])
dp_model = model
update_lr_flag = True
# Assure in training mode
dp_model.train()
for name, param in model.named_parameters():
print(name)
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
CE_ac = utils.CE_ac()
optim_para = model.parameters()
optimizer = utils.build_optimizer(optim_para, opt)
# Load the optimizer
if vars(opt).get('start_from', None) is not None and os.path.isfile(
os.path.join(opt.checkpoint_path, 'optimizer' + opt.id + format(int(opt.start_from),'04')+'.pth')):
optimizer.load_state_dict(torch.load(os.path.join(
opt.checkpoint_path, 'optimizer' + opt.id + format(int(opt.start_from),'04')+'.pth')))
optimizer.zero_grad()
accumulate_iter = 0
train_loss = 0
reward = np.zeros([1,1])
sim_lambda = opt.sim_lambda
reset_optimzer_index = 1
while True:
if opt.self_critical_after != -1 and epoch >= opt.self_critical_after and reset_optimzer_index :
opt.learning_rate_decay_start = opt.self_critical_after
opt.learning_rate_decay_rate = opt.learning_rate_decay_rate_rl
opt.learning_rate = opt.learning_rate_rl
reset_optimzer_index = 0
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_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(opt.train_split)
print('Read data:', time.time() - start)
torch.cuda.synchronize()
start = time.time()
tmp = [data['labels'], data['masks'], data['mods']]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
labels, masks, mods = tmp
tmp = [data['att_feats'], data['att_masks'], data['attr_feats'], data['attr_masks'],data['rela_feats'], data['rela_masks']]
tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
att_feats, att_masks, attr_feats, attr_masks, rela_feats, rela_masks = tmp
rs_data = {}
rs_data['att_feats'] = att_feats
rs_data['att_masks'] = att_masks
rs_data['attr_feats'] = attr_feats
rs_data['attr_masks'] = attr_masks
rs_data['rela_feats'] = rela_feats
rs_data['rela_masks'] = rela_masks
if not sc_flag:
logits, cw_logits = dp_model(rs_data, labels)
ac = CE_ac(logits,labels[:,1:], masks[:,1:])
print('ac :{0}'.format(ac))
loss_lan = crit(logits,labels[:,1:], masks[:,1:])
else:
gen_result, sample_logprobs, cw_logits = dp_model(rs_data,
opt={'sample_max':0}, mode='sample')
reward = get_self_critical_reward(dp_model, rs_data, data, gen_result, opt)
loss_lan = rl_crit(sample_logprobs, gen_result.data, torch.from_numpy(reward).float().cuda())
loss_cw = crit(cw_logits, mods[:, 1:], masks[:, 1:])
ac2 = CE_ac(cw_logits, mods[:, 1:], masks[:, 1:])
print('ac :{0}'.format(ac2))
if epoch < opt.step2_train_after:
loss = loss_lan + sim_lambda*loss_cw
else:
loss = loss_lan
accumulate_iter = accumulate_iter + 1
loss = loss/opt.accumulate_number
loss.backward()
if accumulate_iter % opt.accumulate_number == 0:
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
optimizer.zero_grad()
iteration += 1
accumulate_iter = 0
train_loss = loss.item()*opt.accumulate_number
train_loss_lan = loss_lan.item()
train_loss_cw = loss_cw.item()
end = time.time()
if not sc_flag:
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, end - start))
print("train_loss_lan = {:.3f}, train_loss_cw = {:.3f}" \
.format(train_loss_lan, train_loss_cw))
else:
print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(reward[:, 0]), end - start))
print("train_loss_lan = {:.3f}, train_loss_cw = {:.3f}" \
.format(train_loss_lan, train_loss_cw))
print('lr:{0}'.format(opt.current_lr))
torch.cuda.synchronize()
# Update the iteration and epoch
if data['bounds']['wrapped']:
epoch += 1
update_lr_flag = True
# Write the training loss summary
if (iteration % opt.losses_log_every == 0) and (accumulate_iter % opt.accumulate_number == 0):
add_summary_value(tb_summary_writer, 'train_loss', train_loss, iteration)
add_summary_value(tb_summary_writer, 'train_loss_lan', train_loss_lan, iteration)
add_summary_value(tb_summary_writer, 'train_loss_cw', train_loss_cw, iteration)
add_summary_value(tb_summary_writer, 'learning_rate', opt.current_lr, iteration)
add_summary_value(tb_summary_writer, 'scheduled_sampling_prob', model.ss_prob, iteration)
add_summary_value(tb_summary_writer, 'ac', ac, iteration)
if sc_flag:
add_summary_value(tb_summary_writer, 'avg_reward', np.mean(reward[:,0]), iteration)
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) and (accumulate_iter % opt.accumulate_number == 0):
# eval model
eval_kwargs = {'split': 'test',
'dataset': opt.input_json}
eval_kwargs.update(vars(opt))
#val_loss, predictions, lang_stats = eval_utils_rs3.eval_split(dp_model, crit, loader, eval_kwargs)
# Write validation result into summary
# add_summary_value(tb_summary_writer, 'validation loss', val_loss, iteration)
# if lang_stats is not None:
# for k,v in lang_stats.items():
# add_summary_value(tb_summary_writer, k, v, iteration)
# 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
current_score=0
best_flag = False
if True: # if true
save_id = iteration/opt.save_checkpoint_every
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'+opt.id+format(int(save_id),'04')+'.pth')
torch.save(model.state_dict(), checkpoint_path)
print("model saved to {}".format(checkpoint_path))
optimizer_path = os.path.join(opt.checkpoint_path, 'optimizer'+opt.id+format(int(save_id),'04')+'.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+format(int(save_id),'04')+'.pkl'), 'wb') as f:
cPickle.dump(infos, f)
with open(os.path.join(opt.checkpoint_path, 'histories_'+opt.id+format(int(save_id),'04')+'.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(opt):
opt.use_att = utils.if_use_att(opt)
loader = DataLoader(opt)
opt.vocab_size = loader.vocab_size
opt.seq_length = loader.seq_length
tf_summary_writer = tf and SummaryWriter(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)
best_val_score_vse = infos.get('best_val_score_vse', None)
model = models.JointModel(opt)
model.cuda()
update_lr_flag = True
# Assure in training mode
model.train()
optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
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')):
state_dict = torch.load(os.path.join(opt.start_from, 'optimizer.pth'))
if len(state_dict['state']) == len(optimizer.state_dict()['state']):
optimizer.load_state_dict(state_dict)
else:
print(
'Optimizer param group number not matched? There must be new parameters. Reinit the optimizer.'
)
init_scorer(opt.cached_tokens)
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
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.learning_rate
# 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.caption_generator.ss_prob = opt.ss_prob
# Assign retrieval loss weight
if epoch > opt.retrieval_reward_weight_decay_start and opt.retrieval_reward_weight_decay_start >= 0:
frac = (epoch - opt.retrieval_reward_weight_decay_start
) // opt.retrieval_reward_weight_decay_every
model.retrieval_reward_weight = opt.retrieval_reward_weight * (
opt.retrieval_reward_weight_decay_rate**frac)
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['att_masks'],
data['labels'], data['masks']
]
tmp = utils.var_wrapper(tmp)
fc_feats, att_feats, att_masks, labels, masks = tmp
optimizer.zero_grad()
loss = model(fc_feats, att_feats, att_masks, labels, masks, data)
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, end - start))
prt_str = ""
for k, v in model.loss().items():
prt_str += "{} = {:.3f} ".format(k, v)
print(prt_str)
# 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:
tf_summary_writer.add_scalar('train_loss', train_loss,
iteration)
for k, v in model.loss().items():
tf_summary_writer.add_scalar(k, v, iteration)
tf_summary_writer.add_scalar('learning_rate', opt.current_lr,
iteration)
tf_summary_writer.add_scalar('scheduled_sampling_prob',
model.caption_generator.ss_prob,
iteration)
tf_summary_writer.add_scalar('retrieval_reward_weight',
model.retrieval_reward_weight,
iteration)
tf_summary_writer.file_writer.flush()
loss_history[iteration] = train_loss
lr_history[iteration] = opt.current_lr
ss_prob_history[iteration] = model.caption_generator.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))
# Load the retrieval model for evaluation
val_loss, predictions, lang_stats = eval_utils.eval_split(
model, loader, eval_kwargs)
# Write validation result into summary
if tf is not None:
for k, v in val_loss.items():
tf_summary_writer.add_scalar('validation ' + k, v,
iteration)
for k, v in lang_stats.items():
tf_summary_writer.add_scalar(k, v, iteration)
tf_summary_writer.add_text(
'Captions',
'.\n\n'.join([_['caption'] for _ in predictions[:100]]),
iteration)
#tf_summary_writer.add_image('images', utils.make_summary_image(), iteration)
#utils.make_html(opt.id, iteration)
tf_summary_writer.file_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['SPICE'] * 100
else:
current_score = -val_loss['loss_cap']
current_score_vse = val_loss.get(opt.vse_eval_criterion, 0) * 100
best_flag = False
best_flag_vse = 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
if best_val_score_vse is None or current_score_vse > best_val_score_vse:
best_val_score_vse = current_score_vse
best_flag_vse = 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))
checkpoint_path = os.path.join(opt.checkpoint_path,
'model-%d.pth' % (iteration))
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['best_val_score_vse'] = best_val_score_vse
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,
'infos_' + opt.id + '-%d.pkl' % (iteration)),
'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)
if best_flag_vse:
checkpoint_path = os.path.join(opt.checkpoint_path,
'model_vse-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_vse_' + 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(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()
logger = Logger(opt)
update_lr_flag = True
# Assure in training mode
model.train()
crit = utils.LanguageModelCriterion()
optimizer_G = 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_G.load_state_dict(
torch.load(os.path.join(opt.start_from, 'optimizer.pth')))
while True:
if update_lr_flag:
opt, sc_flag, update_lr_flag, model, optimizer_G = update_lr(
opt, epoch, model, optimizer_G)
# Load data from train split (0)
data = loader.get_batch('train', seq_per_img=opt.seq_per_img)
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_G.zero_grad()
if not sc_flag:
loss = crit(model(fc_feats, att_feats, labels), labels[:, 1:],
masks[:, 1:])
loss.backward()
else:
pass
utils.clip_gradient(optimizer_G, opt.grad_clip)
optimizer_G.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
log = "iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, loss.data.cpu().numpy()[0], end - start)
logger.write(log)
# 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)
tf_summary_writer.flush()
loss_history[iteration] = train_loss
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, logger, eval_kwargs)
logger.write_dict(lang_stats)
# Write validation result into summary
if tf is not None:
add_summary_value(tf_summary_writer, 'validation loss',
val_loss, iteration)
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_G.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(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()
#model_D = Discriminator(opt)
#model_D.load_state_dict(torch.load('save/model_D.pth'))
#model_D.cuda()
#criterion_D = nn.CrossEntropyLoss(size_average=True)
model_E = Distance(opt)
model_E.load_state_dict(
torch.load('save/model_E_NCE/model_E_10epoch.pthsfdasdfadf'))
model_E.cuda()
criterion_E = nn.CosineEmbeddingLoss(margin=0, size_average=True)
#criterion_E = nn.CosineSimilarity()
logger = Logger(opt)
update_lr_flag = True
# Assure in training mode
model.train()
#model_D.train()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer_G = optim.Adam(model.parameters(),
lr=opt.learning_rate,
weight_decay=opt.weight_decay)
#optimizer_D = optim.Adam(model_D.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_G.load_state_dict(
torch.load(os.path.join(opt.start_from, 'optimizer.pth')))
while True:
if update_lr_flag:
opt, sc_flag, update_lr_flag, model, optimizer_G = update_lr(
opt, epoch, model, optimizer_G)
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 = [data['fc_feats'], data['labels'], data['masks']]
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
#fc_feats, att_feats, labels, masks = tmp
fc_feats, labels, masks = tmp
############################################################################################################
############################################ REINFORCE TRAINING ############################################
############################################################################################################
if 1: #iteration % opt.D_scheduling != 0:
optimizer_G.zero_grad()
if not sc_flag:
loss = crit(model(fc_feats, labels), labels[:, 1:], masks[:,
1:])
else:
gen_result, sample_logprobs = model.sample(
fc_feats, {'sample_max': 0})
#reward = get_self_critical_reward(model, fc_feats, att_feats, data, gen_result)
sc_reward = get_self_critical_reward(model, fc_feats, data,
gen_result, logger)
#gan_reward = get_gan_reward(model, model_D, criterion_D, fc_feats, data, logger) # Criterion_D = nn.XEloss()
distance_loss_reward1 = get_distance_reward(
model,
model_E,
criterion_E,
fc_feats,
data,
logger,
is_mismatched=False) # criterion_E = nn.CosEmbedLoss()
distance_loss_reward2 = get_distance_reward(
model,
model_E,
criterion_E,
fc_feats,
data,
logger,
is_mismatched=True) # criterion_E = nn.CosEmbedLoss()
#cosine_reward = get_distance_reward(model, model_E, criterion_E, fc_feats, data, logger) # criterion_E = nn.CosSim()
reward = distance_loss_reward1 + distance_loss_reward2
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(reward).float().cuda(),
requires_grad=False))
loss.backward()
utils.clip_gradient(optimizer_G, opt.grad_clip)
optimizer_G.step()
train_loss = loss.data[0]
torch.cuda.synchronize()
end = time.time()
if not sc_flag:
log = "iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, end - start)
logger.write(log)
else:
log = "iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(reward[:,0]), end - start)
logger.write(log)
######################################################################################################
############################################ GAN TRAINING ############################################
######################################################################################################
else: #elif iteration % opt.D_scheduling == 0: # gan training
model_D.zero_grad()
optimizer_D.zero_grad()
fc_feats_temp = Variable(fc_feats.data.cpu(), volatile=True).cuda()
labels = Variable(labels.data.cpu()).cuda()
sample_res, sample_logprobs = model.sample(
fc_feats_temp, {'sample_max': 0}) #640, 16
greedy_res, greedy_logprobs = model.sample(
fc_feats_temp, {'sample_max': 1}) #640, 16
gt_res = labels # 640, 18
sample_res_embed = model.embed(Variable(sample_res))
greedy_res_embed = model.embed(Variable(greedy_res))
gt_res_embed = model.embed(gt_res)
f_label = Variable(
torch.FloatTensor(data['fc_feats'].shape[0]).cuda())
r_label = Variable(
torch.FloatTensor(data['fc_feats'].shape[0]).cuda())
f_label.data.fill_(0)
r_label.data.fill_(1)
f_D_output = model_D(sample_res_embed.detach(), fc_feats.detach())
f_loss = criterion_D(f_D_output, f_label.long())
f_loss.backward()
r_D_output = model_D(gt_res_embed.detach(), fc_feats.detach())
r_loss = criterion_D(r_D_output, r_label.long())
r_loss.backward()
D_loss = f_loss + r_loss
optimizer_D.step()
torch.cuda.synchronize()
log = 'iter {} (epoch {}), Discriminator loss : {}'.format(
iteration, epoch,
D_loss.data.cpu().numpy()[0])
logger.write(log)
# 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, logger, eval_kwargs)
logger.write_dict(lang_stats)
# Write validation result into summary
if tf is not None:
add_summary_value(tf_summary_writer, 'validation loss',
val_loss, iteration)
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_G.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
