def save_checkpoint(model, infos, optimizer, histories=None, append=''):
if len(append) > 0:
append = '_' + append
# if checkpoint_path doesn't exist
if not os.path.isdir(opt.checkpoint_path):
os.makedirs(opt.checkpoint_path)
checkpoint_path = os.path.join(opt.checkpoint_path,
'model%s.pth' % (append))
torch.save(model.state_dict(), checkpoint_path)
print("Save model state to {}".format(checkpoint_path))
optimizer_path = os.path.join(opt.checkpoint_path,
'optimizer%s.pth' % (append))
torch.save(optimizer.state_dict(), optimizer_path)
print("Save model optimizer to {}".format(optimizer_path))
with open(
os.path.join(opt.checkpoint_path,
'infos' + '%s.pkl' % (append)), 'wb') as f:
utils.pickle_dump(infos, f)
print("Save training information to {}".format(
os.path.join(opt.checkpoint_path,
'infos' + '%s.pkl' % (append))))
if histories:
with open(
os.path.join(opt.checkpoint_path,
'histories' + '%s.pkl' % (append)),
'wb') as f:
utils.pickle_dump(histories, f)
print("Save training historyes to {}".format(
os.path.join(opt.checkpoint_path,
'histories' + '%s.pkl' % (append))))
def main(params):
imgs = json.load(open(params['input_json'], 'r'))
dict_json = json.load(open(params['dict_json'], 'r'))
itow = dict_json['ix_to_word']
wtoi = {w: i for i, w in itow.items()}
# Load bpe
if 'bpe' in dict_json:
import tempfile
import codecs
codes_f = tempfile.NamedTemporaryFile(delete=False)
codes_f.close()
with open(codes_f.name, 'w') as f:
f.write(dict_json['bpe'])
with codecs.open(codes_f.name, encoding='UTF-8') as codes:
bpe = apply_bpe.BPE(codes)
params.bpe = bpe
imgs = imgs['images']
ngram_words, ngram_idxs, ref_len = build_dict(imgs, wtoi, params)
utils.pickle_dump({
'document_frequency': ngram_words,
'ref_len': ref_len
}, open(params['output_pkl'] + '-words.p', 'wb'))
utils.pickle_dump({
'document_frequency': ngram_idxs,
'ref_len': ref_len
}, open(params['output_pkl'] + '-idxs.p', 'wb'))
def save_checkpoint(model, cnn_model, infos, optimizer, cnn_optimizer, histories=None, append=''):
if len(append) > 0:
append = '-' + append
# if checkpoint_path doesn't exist
if not os.path.isdir(opt.checkpoint_path):
os.makedirs(opt.checkpoint_path)
checkpoint_path = os.path.join(
opt.checkpoint_path, 'model%s.pth' % (append))
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%s.pth' % (append))
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%s.pth' % (append))
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%s-cnn.pth' % (append))
torch.save(cnn_optimizer.state_dict(), cnn_optimizer_path)
with open(os.path.join(opt.checkpoint_path, 'infos_' + opt.id + '%s.pkl' % (append)), 'wb') as f:
utils.pickle_dump(infos, f)
if histories:
with open(os.path.join(opt.checkpoint_path, 'histories_' + opt.id + '%s.pkl' % (append)), 'wb') as f:
utils.pickle_dump(histories, f)
def language_eval(dataset, preds, model_id, split):
import sys
sys.path.append("coco-caption")
# if 'coco' in dataset:
# annFile = 'coco-caption/annotations/captions_val2014.json'
# elif 'flickr30k' in dataset or 'f30k' in dataset:
# annFile = 'coco-caption/f30k_captions4eval.json'
annFile = 'C:\\Users\\anke\\PycharmProjects\\pythonProject\\scripts\\wiki_captions4eval.json'
from cococaption.pycocotools.coco import COCO
from cococaption.pycocoevalcap.eval import COCOEvalCap
# encoder.FLOAT_REPR = lambda o: format(o, '.3f')
if not os.path.isdir('eval_results'):
os.mkdir('eval_results')
cache_path = os.path.join('eval_results/',
'.cache_' + model_id + '_' + split + '.json')
coco = COCO(annFile)
valids = coco.getImgIds()
# filter results to only those in MSCOCO validation set (will be about a third)
# preds_filt = [p for p in preds if p['image_id'] in valids] uncomment for not-folder test
preds_filt = preds #comment for non-folder test
print('using %d/%d predictions' % (len(preds_filt), len(preds)))
with open(cache_path, 'wb') as f:
utils.pickle_dump(
preds_filt,
f) # serialize to temporary json file. Sigh, COCO API...
cocoRes = coco.loadRes(cache_path)
cocoEval = COCOEvalCap(coco, cocoRes)
cocoEval.params['image_id'] = cocoRes.getImgIds()
cocoEval.evaluate()
# create output dictionary
out = {}
for metric, score in cocoEval.eval.items():
out[metric] = score
imgToEval = cocoEval.imgToEval
for p in preds_filt:
image_id, caption = p['image_id'], p['caption']
imgToEval[image_id]['caption'] = caption
out['bad_count_rate'] = sum([count_bad(_['caption'])
for _ in preds_filt]) / float(len(preds_filt))
outfile_path = os.path.join('eval_results/',
model_id + '_' + split + '.json')
with open(outfile_path, 'w') as outfile:
json.dump({'overall': out, 'imgToEval': imgToEval}, outfile)
return out
def save_checkpoint(model, infos, optimizer, histories=None, append=''):
if len(append) > 0:
append = '-' + append
if not os.path.isdir(opt.checkpoint_path):
os.makedirs(opt.checkpoint_path)
checkpoint_path = os.path.join(opt.checkpoint_path, 'model%s.pth' %(append))
torch.save(model.state_dict(), checkpoint_path)
print("model saved to {}".format(checkpoint_path))
optimizer_path = os.path.join(opt.checkpoint_path, 'optimizer%s.pth' %(append))
torch.save(optimizer.state_dict(), optimizer_path)
with open(os.path.join(opt.checkpoint_path, 'infos%s.pkl' %(append)), 'wb') as f:
utils.pickle_dump(infos, f)
if histories:
with open(os.path.join(opt.checkpoint_path, 'histories%s.pkl' %(append)), 'wb') as f:
utils.pickle_dump(histories, f)
def main(params):
imgs = json.load(open(params['input_json'], 'r'))
itow = json.load(open(params['dict_json'], 'r'))['ix_to_word']
wtoi = {w: i for i, w in itow.items()}
imgs = imgs['images']
ngram_words, ngram_idxs, ref_len = build_dict(imgs, wtoi, params)
utils.pickle_dump({
'document_frequency': ngram_words,
'ref_len': ref_len
}, open(params['output_pkl'] + '-words.p', 'w'))
utils.pickle_dump({
'document_frequency': ngram_idxs,
'ref_len': ref_len
}, open(params['output_pkl'] + '-idxs.p', 'w'))
def save_checkpoint(model, cnn_model, infos, optimizer, histories=None, append=''):
if len(append) > 0:
append = '-' + append
# if checkpoint_path doesn't exist
if not os.path.isdir(opt.checkpoint_path):
os.makedirs(opt.checkpoint_path)
#Transformer model
checkpoint_path = os.path.join(opt.checkpoint_path, 'model%s.pth' % (append))
torch.save(model.state_dict(), checkpoint_path)
print("model saved to {}".format(checkpoint_path))
#CNN model
checkpoint_path = os.path.join(opt.checkpoint_path, 'cnn_model%s.pth' % (append))
if not os.path.exists(checkpoint_path):
torch.save(cnn_model.state_dict(), checkpoint_path)
print("model saved to {}".format(checkpoint_path))
optimizer_path = os.path.join(opt.checkpoint_path, 'optimizer%s.pth' % (append))
torch.save(optimizer.state_dict(), optimizer_path)
with open(os.path.join(opt.checkpoint_path, 'infos_' + opt.id + '%s.pkl' % (append)), 'wb') as f:
utils.pickle_dump(infos, f)
if histories:
with open(os.path.join(opt.checkpoint_path, 'histories_' + opt.id + '%s.pkl' % (append)), 'wb') as f:
utils.pickle_dump(histories, f)
def main(params):
imgs = json.load(open(params['input_json'], 'r'))
# for iccv 19:cnm
dict_json = json.load(open(params['dict_json'], 'r'))
itow = dict_json['ix_to_word']
# for cvpr 19:sgae
# sg_dict_info = np.load(params['dict_json'])['spice_dict'][()]
# itow = sg_dict_info['ix_to_word']
wtoi = {w: i for i, w in itow.items()}
# Load bpe
# if 'bpe' in dict_json:
# import tempfile
# import codecs
# codes_f = tempfile.NamedTemporaryFile(delete=False)
# codes_f.close()
# with open(codes_f.name, 'w') as f:
# f.write(dict_json['bpe'])
# with codecs.open(codes_f.name, encoding='UTF-8') as codes:
# bpe = apply_bpe.BPE(codes)
# params.bpe = bpe
imgs = imgs['images']
ngram_words, ngram_idxs, ref_len = build_dict(imgs, wtoi, params)
utils.pickle_dump({
'document_frequency': ngram_words,
'ref_len': ref_len
}, open(params['output_pkl'] + '-words.p', 'w'))
utils.pickle_dump({
'document_frequency': ngram_idxs,
'ref_len': ref_len
}, open(params['output_pkl'] + '-idxs.p', 'w'))
def main(params):
imgs = json.load(open(params['input_json'], 'r'))
itow = json.load(open(params['dict_json'], 'r'))['ix_to_word']
wtoi = {w: i for i, w in itow.items()}
imgs = imgs['images']
ngram_words, ngram_idxs, ref_len = build_dict(imgs, wtoi, params)
print(cPickle.HIGHEST_PROTOCOL)
utils.pickle_dump(({
'document_frequency': ngram_words,
'ref_len': ref_len
}),
open(params['output_pkl'] + '-words.p', 'wb+'),
protocol=cPickle.HIGHEST_PROTOCOL)
utils.pickle_dump(({
'document_frequency': ngram_idxs,
'ref_len': ref_len
}),
open(params['output_pkl'] + '-idxs.p', 'wb+'),
protocol=cPickle.HIGHEST_PROTOCOL)
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
crit = utils.LanguageModelCriterion()
# Create the Data Loader instance
if opt.image_folder is None or len(opt.image_folder) == 0:
loader = DataLoader(opt)
else:
loader = DataLoaderRaw({
'folder_path': opt.image_folder,
'coco_json': opt.coco_json,
'batch_size': opt.batch_size,
'cnn_model': opt.cnn_model
})
# When eval using provided pretrained model, the vocab may be different from what you have in your cocotalk.json
# So make sure to use the vocab in infos file.
loader.ix_to_word = infos['vocab']
# Set sample options
opt.datset = opt.input_json
loss, split_predictions, lang_stats = eval_utils.eval_split(
model, crit, loader, vars(opt))
print('loss: ', loss)
if lang_stats:
print(lang_stats)
if opt.dump_json == 1:
# dump the json
evout = opt.vis_dump_path
with open(evout, 'wb') as f:
utils.pickle_dump(split_predictions, f)
