def __init__(self, model, opt):
super(LossWrapper, self).__init__()
self.opt = opt
self.model = model
self.crit = utils.LanguageModelCriterion()
self.rl_crit = utils.RewardCriterion()
self.struc_crit = utils.StructureLosses(opt)
def main(opt):
opt_test = opt
test_dataset = VideoDataset(opt_test, 'test')
opt_test["vocab_size"] = test_dataset.get_vocab_size()
opt_test["seq_length"] = test_dataset.max_len
dataset = VideoDataset(opt, 'train')
dataloader = DataLoader(dataset,
batch_size=opt["batch_size"],
shuffle=True)
opt["vocab_size"] = dataset.get_vocab_size()
if opt["model"] == 'S2VTModel':
model = S2VTModel(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"],
opt["dim_word"],
opt['dim_vid'],
rnn_cell=opt['rnn_type'],
n_layers=opt['num_layers'],
rnn_dropout_p=opt["rnn_dropout_p"])
elif opt["model"] == "S2VTAttModel":
encoder = EncoderRNN(
opt["dim_vid"],
opt["dim_hidden"],
# bidirectional=opt["bidirectional"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"])
second_lstm = Two_Lstm(
opt["dim_vid"],
opt["dim_hidden"],
# bidirectional=opt["bidirectional"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"])
decoder = DecoderRNN(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"],
opt["dim_word"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"])
# bidirectional=opt["bidirectional"])
model = S2VTAttModel(encoder, second_lstm, decoder)
model = model.cuda()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = optim.Adam(model.parameters(),
lr=opt["learning_rate"],
weight_decay=opt["weight_decay"])
exp_lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=opt["learning_rate_decay_every"],
gamma=opt["learning_rate_decay_rate"])
train(dataloader, model, crit, optimizer, exp_lr_scheduler, opt, rl_crit,
opt_test, test_dataset)
def __init__(self, model, opt):
super(LossWrapper, self).__init__()
self.opt = opt
self.model = model
if opt.label_smoothing > 0:
self.crit = utils.LabelSmoothing(smoothing=opt.label_smoothing)
else:
self.crit = utils.LanguageModelCriterion()
self.rl_crit = utils.RewardCriterion()
self.ppo_crit = utils.PPORewardCriterion() ## Added this for ppo 9/sep/2019
self.old_sample_logprobs = torch.zeros(50,16).to('cuda')
def main(opt):
dataset = VideoDataset(opt, 'train')
dataloader = DataLoader(dataset,
batch_size=opt["batch_size"],
num_workers=8,
shuffle=True)
opt["vocab_size"] = dataset.get_vocab_size()
if opt["model"] == 'S2VTModel':
model = S2VTModel(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"],
opt["dim_word"],
opt['dim_vid'],
rnn_cell=opt['rnn_type'],
n_layers=opt['num_layers'],
bidirectional=opt["bidirectional"],
rnn_dropout_p=opt["rnn_dropout_p"]).cuda()
elif opt["model"] == "S2VTAttModel":
encoder = EncoderRNN(opt["dim_vid"],
opt["dim_hidden"],
n_layers=opt['num_layers'],
bidirectional=opt["bidirectional"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"])
decoder = DecoderRNN(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"],
opt["dim_word"],
n_layers=opt['num_layers'],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"],
bidirectional=opt["bidirectional"])
model = S2VTAttModel(encoder, decoder).cuda()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = optim.Adam(model.parameters(),
lr=opt["learning_rate"],
weight_decay=opt["weight_decay"])
exp_lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=opt["learning_rate_decay_every"],
gamma=opt["learning_rate_decay_rate"])
model.load_state_dict(
torch.load(
"C:\\Users\\Shumpu\\VideoCaptioningAttack\\video_caption_pytorch\\save\\vgg16_model_460.pth"
))
train(dataloader, model, crit, optimizer, exp_lr_scheduler, opt, rl_crit)
def __init__(self, model, opt): super(LossWrapper, self).__init__() self.opt = opt self.model = model if opt.label_smoothing > 0: self.crit = utils.LabelSmoothing(smoothing=opt.label_smoothing) else: self.crit = utils.LanguageModelCriterion() self.rl_crit = utils.RewardCriterion() if opt.att_supervise: if opt.att_sup_crit == 'KL': self.kl_crit=nn.KLDivLoss(reduction='batchmean') elif opt.att_sup_crit == 'NLL': self.nll = nn.NLLLoss() elif opt.att_sup_crit == 'ExtendNLL': self.extendnll = utils.ExtendNLLCrit() else: raise NotImplementedError self.min_value=1e-8
def main(opt):
train_dataset = VideoDataset(opt, 'train')
train_dataloader = DataLoader(train_dataset,
batch_size=opt.batch_size,
shuffle=True)
opt.vocab_size = train_dataset.vocab_size
opt.seq_length = train_dataset.seq_length
val_dataset = VideoDataset(opt, 'val')
val_dataloader = DataLoader(val_dataset,
batch_size=opt.batch_size,
shuffle=True)
if opt.model == 'S2VTModel':
model = S2VTModel(opt.vocab_size,
opt.seq_length,
opt.dim_hidden,
opt.dim_word,
rnn_dropout_p=opt.rnn_dropout_p).cuda()
elif opt.model == "Vid2seq":
encoder = EncoderRNN(opt.dim_vid, opt.dim_hidden)
decoder = DecoderRNN(opt.vocab_size,
opt.seq_length,
opt.dim_hidden,
use_attention=True,
rnn_dropout_p=opt.rnn_dropout_p)
model = Vid2seq(encoder, decoder).cuda()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = optim.Adam(model.parameters(),
lr=opt.learning_rate,
weight_decay=opt.weight_decay)
exp_lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=opt.learning_rate_decay_every,
gamma=opt.learning_rate_decay_rate)
if not os.path.isdir(opt.checkpoint_path):
os.mkdir(opt.checkpoint_path)
train(train_dataloader, val_dataloader, model, crit, optimizer,
exp_lr_scheduler, opt, rl_crit)
def __init__(self, model, opt):
super(LossWrapper, self).__init__()
self.opt = opt
self.model = model
if opt.label_smoothing > 0:
self.crit = utils.LabelSmoothing(smoothing=opt.label_smoothing)
else:
self.crit = utils.LanguageModelCriterion()
self.rl_crit = utils.RewardCriterion()
self.struc_crit = utils.StructureLosses(opt)
if opt.vse_model != 'None':
self.vse = VSEFCModel(opt)
for p in self.vse.parameters():
p.requires_grad = False
self.retrieval_reward_weight = opt.retrieval_reward_weight #
self.vse.load_state_dict({
k[4:]: v
for k, v in torch.load(opt.initialize_retrieval).items()
if 'vse.' in k
})
self.retrieval_reward_weight = 0
def main(opt):
dataset = VideoDataset(opt, 'train')
dataloader = DataLoader(
dataset, batch_size=opt["batch_size"], shuffle=True)
opt["vocab_size"] = dataset.get_vocab_size()
if opt["model"] == 'S2VTModel':
model = S2VTModel(opt["vocab_size"], opt["max_len"], opt["dim_hidden"], opt["dim_word"],
rnn_dropout_p=opt["rnn_dropout_p"]).cuda()
elif opt["model"] == "S2VTAttModel":
encoder = EncoderRNN(opt["dim_vid"], opt["dim_hidden"], bidirectional=opt["bidirectional"],
input_dropout_p=opt["input_dropout_p"], rnn_dropout_p=opt["rnn_dropout_p"])
decoder = DecoderRNN(opt["vocab_size"], opt["max_len"], opt["dim_hidden"], opt["dim_word"],
input_dropout_p=opt["input_dropout_p"],
rnn_dropout_p=opt["rnn_dropout_p"], bidirectional=opt["bidirectional"])
model = S2VTAttModel(encoder, decoder).cuda()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = optim.Adam(
model.parameters(), lr=opt["learning_rate"], weight_decay=opt["weight_decay"])
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt["learning_rate_decay_every"],
gamma=opt["learning_rate_decay_rate"])
train(dataloader, model, crit, optimizer, exp_lr_scheduler, opt, rl_crit)
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
def train(opt):
# Deal with feature things before anything
opt.use_fc, opt.use_att = utils.if_use_feat(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
with open(os.path.join(opt.start_from, 'infos_' + opt.id + '.pkl'),
'rb') as f:
infos = utils.pickle_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 = utils.pickle_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)
epoch_done = True
# Assure in training mode
dp_model.train()
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', 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')))
total_loss = 0
times = 0
while True:
if epoch_done:
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)
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
times += 1
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()
total_loss = total_loss + train_loss
torch.cuda.synchronize()
end = time.time()
if not sc_flag:
print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, train_loss, end - start))
else:
print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \
.format(iteration, epoch, np.mean(reward[:,0]), end - start))
# Update the iteration and epoch
iteration += 1
if data['bounds']['wrapped']:
# epoch += 1
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.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):
if data['bounds']['wrapped']:
epoch += 1
# eval model
eval_kwargs = {
'split': 'val',
'dataset': opt.input_json,
'verbose': False
}
eval_kwargs.update(vars(opt))
val_loss, predictions, lang_stats = eval_utils.eval_split(
dp_model, crit, loader, eval_kwargs)
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)
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
f = open('train_log_%s.txt' % opt.id, 'a')
f.write(
'Epoch {}: | Date: {} | TrainLoss: {} | ValLoss: {} | Score: {}'
.format(epoch, str(datetime.now()),
str(total_loss / times), str(val_loss),
str(current_score)))
f.write('\n')
f.close()
print('-------------------wrote to log file')
total_loss = 0
times = 0
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
if not os.path.isdir(opt.checkpoint_path):
os.mkdir(opt.checkpoint_path)
checkpoint_path = os.path.join(opt.checkpoint_path,
'model.pth')
torch.save(model.state_dict(), checkpoint_path)
# print(str(infos['best_val_score']))
print("model saved to {}".format(checkpoint_path))
if opt.save_history_ckpt:
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['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:
utils.pickle_dump(infos, f)
if opt.save_history_ckpt:
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:
utils.pickle_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:
utils.pickle_dump(infos, f)
# Stop if reaching max epochs
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
def train(rank, model, opt, optimizer=None):
torch.manual_seed(opt.seed + rank)
if opt.use_cuda:
torch.cuda.manual_seed(opt.seed + rank)
loader = DataLoader(opt)
index_2_word = loader.get_vocab()
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.load_model_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
iteration = infos.get('iter', 0)
epoch = infos.get('epoch', 0)
val_result_history = infos.get('val_result_history', {})
loss_history = infos.get('loss_history', {})
lr_history = infos.get('lr_history', {})
ss_prob_history = infos.get('ss_prob_history', {})
sorted_lr = sorted(lr_history.items(), key=operator.itemgetter(1))
if opt.load_lr and len(lr_history) > 0:
opt.optim_rl_lr = sorted_lr[0][1] / opt.optim_rl_lr_ratio
loader.iterators = infos.get('iterators', loader.iterators)
loader.split_image_id = infos.get('split_image_id', loader.split_image_id)
entropy_reg = opt.entropy_reg
best_val_score = 0
if opt.load_best_score == 1:
best_val_score = infos.get('best_val_score', None)
update_lr_flag = True
if opt.caption_model == 'show_tell':
crit = utils.LanguageModelCriterion(opt)
rl_crit = utils.RewardCriterion(opt)
elif opt.caption_model == 'review_net':
crit = utils.ReviewNetCriterion(opt)
rl_crit = utils.ReviewNetRewardCriterion(opt)
elif opt.caption_model == 'recurrent_fusion_model':
crit = utils.ReviewNetEnsembleCriterion(opt)
rl_crit = utils.ReviewNetRewardCriterion(opt)
else:
raise Exception("caption_model not supported: {}".format(
opt.caption_model))
if optimizer is None:
if opt.optim == 'adam':
optimizer = optim.Adam(model.parameters(),
lr=opt.optim_rl_lr,
betas=(opt.optim_adam_beta1,
opt.optim_adam_beta2),
weight_decay=opt.optim_weight_decay)
elif opt.optim == 'rmsprop':
optimizer = optim.RMSprop(model.parameters(),
lr=opt.optim_rl_lr,
momentum=opt.optim_momentum,
alpha=opt.optim_rmsprop_alpha,
weight_decay=opt.weight_decay)
elif opt.optim == 'sgd':
optimizer = optim.SGD(model.parameters(),
lr=opt.optim_rl_lr,
momentum=opt.optim_momentum,
weight_decay=opt.optim_weight_decay)
elif opt.optim == 'adagrad':
optimizer = optim.Adagrad(model.parameters(),
lr=opt.optim_rl_lr,
lr_decay=opt.optim_lr_decay,
weight_decay=opt.optim_weight_decay)
elif opt.optim == 'adadelta':
optimizer = optim.Adadelta(model.parameters(),
rho=opt.optim_rho,
eps=opt.optim_epsilon,
lr=opt.optim_rl_lr,
weight_decay=opt.optim_weight_decay)
else:
raise Exception("optim not supported: {}".format(opt.feature_type))
# Load the optimizer
if opt.load_lr and vars(opt).get(
'start_from', None) is not None and os.path.isfile(
os.path.join(opt.start_from,
'optimizer_' + opt.load_model_id + '.pth')):
optimizer.load_state_dict(
torch.load(
os.path.join(opt.start_from,
'optimizer_' + opt.load_model_id + '.pth')))
utils.set_lr(optimizer, opt.optim_rl_lr)
num_period_best = 0
current_score = 0
while True:
if update_lr_flag:
# Assign the learning rate
if epoch > 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.optim_rl_lr * decay_factor
utils.set_lr(optimizer, opt.current_lr) # set the decayed rate
else:
opt.current_lr = opt.optim_rl_lr
update_lr_flag = False
start = time.time()
data = loader.get_batch('train')
if opt.use_cuda:
torch.cuda.synchronize()
if opt.feature_type == 'feat_array':
fc_feat_array = data['fc_feats_array']
att_feat_array = data['att_feats_array']
assert (len(fc_feat_array) == len(att_feat_array))
for feat_id in range(len(fc_feat_array)):
if opt.use_cuda:
fc_feat_array[feat_id] = Variable(
torch.from_numpy(fc_feat_array[feat_id]),
requires_grad=False).cuda()
att_feat_array[feat_id] = Variable(
torch.from_numpy(att_feat_array[feat_id]),
requires_grad=False).cuda()
else:
fc_feat_array[feat_id] = Variable(torch.from_numpy(
fc_feat_array[feat_id]),
requires_grad=False)
att_feat_array[feat_id] = Variable(torch.from_numpy(
att_feat_array[feat_id]),
requires_grad=False)
tmp = [data['labels'], data['masks'], data['top_words']]
if opt.use_cuda:
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
else:
tmp = [
Variable(torch.from_numpy(_), requires_grad=False)
for _ in tmp
]
labels, masks, top_words = tmp
else:
tmp = [
data['fc_feats'], data['att_feats'], data['labels'],
data['masks'], data['top_words']
]
if opt.use_cuda:
tmp = [
Variable(torch.from_numpy(_), requires_grad=False).cuda()
for _ in tmp
]
else:
tmp = [
Variable(torch.from_numpy(_), requires_grad=False)
for _ in tmp
]
fc_feats, att_feats, labels, masks, top_words = tmp
optimizer.zero_grad()
if opt.caption_model == 'show_tell':
gen_result, sample_logprobs, logprobs_all = model.sample(
fc_feats, att_feats, {'sample_max': 0})
rewards = get_rewards.get_self_critical_reward(
index_2_word, model, fc_feats, att_feats, data, gen_result,
opt)
sample_logprobs_old = Variable(sample_logprobs.data,
requires_grad=False)
if opt.use_cuda:
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(rewards).float().cuda(),
requires_grad=False), logprobs_all, entropy_reg,
sample_logprobs_old, opt)
else:
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(rewards).float(),
requires_grad=False), logprobs_all, entropy_reg,
sample_logprobs_old, opt)
elif opt.caption_model == 'recurrent_fusion_model':
gen_result, sample_logprobs, logprobs_all, top_pred = model.sample(
fc_feat_array, att_feat_array, {'sample_max': 0})
rewards = get_rewards.get_self_critical_reward_feat_array(
index_2_word, model, fc_feat_array, att_feat_array, data,
gen_result, opt)
sample_logprobs_old = Variable(sample_logprobs.data,
requires_grad=False)
if opt.use_cuda:
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(rewards).float().cuda(),
requires_grad=False), logprobs_all, entropy_reg,
top_pred, top_words, opt.reason_weight,
sample_logprobs_old, opt)
else:
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(rewards).float(),
requires_grad=False), logprobs_all, entropy_reg,
top_pred, top_words, opt.reason_weight,
sample_logprobs_old, opt)
elif opt.caption_model == 'review_net':
gen_result, sample_logprobs, logprobs_all, top_pred = model.sample(
fc_feats, att_feats, {'sample_max': 0})
rewards = get_rewards.get_self_critical_reward(
index_2_word, model, fc_feats, att_feats, data, gen_result,
opt)
sample_logprobs_old = Variable(sample_logprobs.data,
requires_grad=False)
if opt.use_cuda:
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(rewards).float().cuda(),
requires_grad=False), logprobs_all, entropy_reg,
top_pred, top_words, opt.reason_weight,
sample_logprobs_old, opt)
else:
loss = rl_crit(
sample_logprobs, gen_result,
Variable(torch.from_numpy(rewards).float(),
requires_grad=False), logprobs_all, entropy_reg,
top_pred, top_words, opt.reason_weight,
sample_logprobs_old, opt)
else:
raise Exception("caption_model not supported: {}".format(
opt.caption_model))
if opt.use_ppo and opt.ppo_k > 0:
loss.backward(retain_graph=True)
else:
loss.backward()
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
train_loss = loss.data[0]
if opt.use_ppo:
for i in range(opt.ppo_k):
print(i)
optimizer.zero_grad()
loss.backward(retain_graph=True)
utils.clip_gradient(optimizer, opt.grad_clip)
optimizer.step()
if opt.use_cuda:
torch.cuda.synchronize()
end = time.time()
# 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:
loss_history[iteration] = np.mean(rewards[:, 0])
lr_history[iteration] = opt.current_lr
# make evaluation on validation set, and save model
if iteration % opt.save_checkpoint_every == 0:
# eval model
eval_kwargs = {
'eval_split': 'val',
'dataset': opt.input_json,
'caption_model': opt.caption_model,
'reason_weight': opt.reason_weight,
'guiding_l1_penality': opt.guiding_l1_penality,
'use_cuda': opt.use_cuda,
'feature_type': opt.feature_type,
'rank': rank
}
eval_kwargs.update(vars(opt))
eval_kwargs['eval_split'] = 'val'
val_loss, predictions, lang_stats = eval_utils.eval_split(
model, crit, loader, eval_kwargs)
# Write validation result into summary
val_result_history[iteration] = {
'loss': val_loss,
'lang_stats': lang_stats,
'predictions': predictions
}
print("iter {} (epoch {}), val_loss = {:.3f}".format(
iteration, epoch, val_loss))
# 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 best_val_score is None or current_score > best_val_score:
best_val_score = current_score
best_flag = True
num_period_best = 1
else:
num_period_best = num_period_best + 1
# Dump miscalleous informations
infos['iter'] = iteration
infos['epoch'] = epoch
infos['iterators'] = loader.iterators
infos['split_image_id'] = loader.split_image_id
infos['best_val_score'] = best_val_score
infos['opt'] = opt
infos['val_result_history'] = val_result_history
infos['loss_history'] = loss_history
infos['lr_history'] = lr_history
infos['ss_prob_history'] = ss_prob_history
infos['vocab'] = loader.get_vocab()
with open(
os.path.join(
opt.checkpoint_path,
'rl_infos_' + opt.id + '_' + str(rank) + '.pkl'),
'wb') as f:
cPickle.dump(infos, f)
if best_flag:
checkpoint_path = os.path.join(
opt.checkpoint_path,
'rl_model_' + opt.id + '_' + str(rank) + '-best.pth')
torch.save(model.state_dict(), checkpoint_path)
optimizer_path = os.path.join(
opt.checkpoint_path,
'rl_optimizer_' + opt.id + '_' + str(rank) + '-best.pth')
torch.save(optimizer.state_dict(), optimizer_path)
print("model saved to {}".format(checkpoint_path))
with open(
os.path.join(
opt.checkpoint_path, 'rl_infos_' + opt.id + '_' +
str(rank) + '-best.pkl'), 'wb') as f:
cPickle.dump(infos, f)
if num_period_best >= opt.num_eval_no_improve:
print('no improvement, exit')
sys.exit()
print("rank {}, iter {}, (epoch {}), avg_reward: {:.3f}, train_loss: {}, learning rate: {}, current cider: {:.3f}, best cider: {:.3f}, time: {:.3f}" \
.format(rank, iteration, epoch, np.mean(rewards[:, 0]), train_loss, opt.current_lr, current_score, best_val_score, (end-start)))
iteration += 1
# Stop if reaching max epochs
if epoch >= opt.max_epochs and opt.max_epochs != -1:
break
def main(opt):
dataset = VideoDataset(opt, 'train')
dataloader = DataLoader(dataset,
batch_size=opt["batch_size"],
shuffle=True)
opt["vocab_size"] = dataset.get_vocab_size()
if opt["model"] == 'S2VTModel':
model = S2VTModel(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"],
opt["dim_word"],
opt['dim_vid'],
rnn_cell=opt['rnn_type'],
n_layers=opt['num_layers'],
rnn_dropout_p=opt["rnn_dropout_p"])
elif opt["model"] == "S2VTAttModel":
encoder = EncoderRNN(opt["dim_vid"],
opt["dim_hidden"],
bidirectional=opt["bidirectional"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"])
# # 声音encoder
# encoder_voice = EncoderRNN(
# opt["dim_voice"],
# opt["dim_hidden"],
# bidirectional=opt["bidirectional"],
# input_dropout_p=opt["input_dropout_p"],
# rnn_cell=opt['rnn_type'],
# rnn_dropout_p=opt["rnn_dropout_p"])
# 手语encoder
if opt['with_hand'] == 1:
encoder_hand = EncoderRNN(opt["dim_hand"],
opt["dim_hand_hidden"],
bidirectional=opt["bidirectional"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"])
decoder = DecoderRNN(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"] + opt["dim_hand_hidden"],
opt["dim_word"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"],
bidirectional=opt["bidirectional"])
model = S2VTAttModel(encoder, encoder_hand, decoder)
else:
decoder = DecoderRNN(opt["vocab_size"],
opt["max_len"],
opt["dim_hidden"],
opt["dim_word"],
input_dropout_p=opt["input_dropout_p"],
rnn_cell=opt['rnn_type'],
rnn_dropout_p=opt["rnn_dropout_p"],
bidirectional=opt["bidirectional"])
model = S2VTAttModel(encoder, None, decoder)
# model = S2VTAttModel(encoder, encoder_voice, encoder_hand, decoder)
model = model.cuda()
crit = utils.LanguageModelCriterion()
rl_crit = utils.RewardCriterion()
optimizer = optim.Adam(model.parameters(),
lr=opt["learning_rate"],
weight_decay=opt["weight_decay"])
exp_lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=opt["learning_rate_decay_every"],
gamma=opt["learning_rate_decay_rate"])
# print(dataloader)
# print(crit)
# print(optimizer)
train(dataloader, model, crit, optimizer, exp_lr_scheduler, opt, rl_crit)
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):
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