def main(args):
# set up model
checkpoint_path = osp.join(args.checkpoint_dir, '%s.pth' % args.id)
checkpoint = torch.load(checkpoint_path)
opt = checkpoint['opt']
model = Navigator(checkpoint['opt'])
model.load_state_dict(checkpoint['model_state'])
model.cuda()
print('model set up.')
# set up criterions
nll_crit = SeqModelCriterion().cuda()
mse_crit = MaskedMSELoss().cuda()
# set up loader
loader_kwargs = {
'data_json': args.data_json,
'data_h5': args.data_h5,
'path_feats_dir': args.path_feats_dir,
'path_images_dir': args.path_images_dir,
'split': args.split,
'max_seq_length': 100,
'requires_imgs': False,
'nav_types': ['room'],
'question_types': ['all'],
}
dataset = NavImitationDataset(**loader_kwargs)
# evaluate
predictions, overall_nll, overall_teacher_forcing_acc, overall_mse, Nav_nll, Nav_teacher_forcing_acc, Nav_cnt = \
evaluate(dataset, model, nll_crit, mse_crit, opt)
# summarize
results_str = 'id[%s] ' % args.id
if opt['use_action']: results_str += '[use action]'
if opt['use_curriculum']: results_str += '[use curriculum]'
if opt['use_next']: results_str += '[use next]'
if opt['use_residual']: results_str += '[use residual]'
results_str += '\nsplit[%s]\n' % args.split
results_str += ' nll_loss: %.3f\n' % overall_nll
results_str += ' teacher-forcing acc (%s): %.2f%%,' % (
len(predictions), overall_teacher_forcing_acc * 100.)
results_str += ' on %s objects: %.2f%%,' % (
Nav_cnt['object'], Nav_teacher_forcing_acc['object'] * 100.)
results_str += ' on %s rooms: %.2f%%\n' % (
Nav_cnt['room'], Nav_teacher_forcing_acc['room'] * 100.)
# save
with open(args.result_json, 'w') as f:
json.dump(predictions, f)
f = open(args.report_txt, 'a')
f.write(results_str)
f.write('\n')
f.close()
def imitation(rank, args, shared_nav_model, counter):
# set up tensorboard
writer = tb.SummaryWriter(args.tb_dir, filename_suffix=str(rank))
# set up cuda device
torch.cuda.set_device(args.gpus.index(args.gpus[rank % len(args.gpus)]))
# set up random seeds
random.seed(args.seed + rank)
np.random.randn(args.seed + rank)
torch.manual_seed(args.seed + rank)
# set up loader
train_loader_kwargs = {
'data_json': args.imitation_data_json,
'data_h5': args.imitation_data_h5,
'path_feats_dir': args.path_feats_dir,
'path_images_dir': args.path_images_dir,
'split': 'train',
'max_seq_length': args.max_seq_length,
'requires_imgs': False,
'nav_types': args.nav_types,
'question_types': ['all'],
}
train_dataset = NavImitationDataset(**train_loader_kwargs)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=1)
print('train_loader set up.')
# set up optimizer on shared_nav_model
# lr = 5e-5
lr = args.learning_rate
optimizer = torch.optim.Adam(shared_nav_model.parameters(),
lr=lr,
betas=(args.optim_alpha, args.optim_beta),
eps=args.optim_epsilon,
weight_decay=args.weight_decay)
# set up models
opt = vars(args)
opt['act_to_ix'] = train_dataset.act_to_ix
opt['num_actions'] = len(opt['act_to_ix'])
model = Navigator(opt)
model.cuda()
print('navigator set up.')
# set up criterions
nll_crit = SeqModelCriterion().cuda()
# -
epoch = 0
iters = 0
# train
while True:
for batch in train_loader:
# sync model
model.load_state_dict(shared_nav_model.state_dict())
model.train()
model.cuda()
# batch = {qid, path_ix, house, id, type, phrase, phrase_emb, ego_feats, next_feats, res_feats,
# action_inputs, action_outputs, action_masks, ego_imgs}
ego_feats = batch['ego_feats'].cuda() # (n, L, 3200)
phrase_embs = batch['phrase_emb'].cuda() # (n, 300)
action_inputs = batch['action_inputs'].cuda() # (n, L)
action_outputs = batch['action_outputs'].cuda() # (n, L)
action_masks = batch['action_masks'].cuda() # (n, L)
# forward
# - logprobs (n, L, #actions)
# - output_feats (n, L, rnn_size)
# - pred_feats (n, L, 3200) or None
logprobs, _, pred_feats, _ = model(ego_feats, phrase_embs,
action_inputs)
nll_loss = nll_crit(logprobs, action_outputs, action_masks)
# backward
optimizer.zero_grad()
nll_loss.backward()
clip_model_gradient(model.parameters(), args.grad_clip)
ensure_shared_grads(model.cpu(), shared_nav_model)
optimizer.step()
if iters % 25 == 0:
print('imitation-r%s(ep%s it%s lr%.2E loss%.4f)' %
(rank, epoch, iters, lr, nll_loss))
# write to tensorboard
writer.add_scalar('imitation_rank/nll_loss',
float(nll_loss.item()), counter.value)
# increate iters
iters += 1
# decay learning rate
if args.lr_decay > 0:
if args.im_learning_rate_decay_start > 0 and iters > args.im_learning_rate_decay_start:
frac = (iters - args.im_learning_rate_decay_start
) / args.im_learning_rate_decay_every
decay_factor = 0.1**frac
lr = args.learning_rate * decay_factor
model_utils.set_lr(optimizer, lr)
epoch += 1
def main(args):
# make output directory
if args.checkpoint_dir is None:
args.checkpoint_dir = 'output/nav_object'
if not osp.isdir(args.checkpoint_dir):
os.makedirs(args.checkpoint_dir)
# set random seed
random.seed(args.seed)
np.random.randn(args.seed)
torch.manual_seed(args.seed)
# set up loaders
train_loader_kwargs = {
'data_json': args.data_json,
'data_h5': args.data_h5,
'path_feats_dir': args.path_feats_dir,
'path_images_dir': args.path_images_dir,
'split': 'train',
'max_seq_length': args.max_seq_length,
'requires_imgs': False,
'nav_types': ['object'],
'question_types': ['all'],
}
train_dataset = NavImitationDataset(**train_loader_kwargs)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=args.num_workers)
val_loader_kwargs = {
'data_json': args.data_json,
'data_h5': args.data_h5,
'path_feats_dir': args.path_feats_dir,
'path_images_dir': args.path_images_dir,
'split': 'val',
'max_seq_length': args.max_seq_length,
'requires_imgs': False,
'nav_types': ['object'],
'question_types': ['all'],
}
val_dataset = NavImitationDataset(**val_loader_kwargs)
# set up models
opt = vars(args)
opt['act_to_ix'] = train_dataset.act_to_ix
opt['num_actions'] = len(opt['act_to_ix'])
model = Navigator(opt)
model.cuda()
print('navigator set up.')
# set up criterions
nll_crit = SeqModelCriterion().cuda()
mse_crit = MaskedMSELoss().cuda()
# set up optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=args.learning_rate,
betas=(args.optim_alpha, args.optim_beta),
eps=args.optim_epsilon,
weight_decay=args.weight_decay)
# resume from checkpoint
infos = {}
iters = infos.get('iters', 0)
epoch = infos.get('epoch', 0)
val_nll_history = infos.get('val_nll_history', {})
val_mse_history = infos.get('val_mse_history', {})
val_teacher_forcing_acc_history = infos.get(
'val_teacher_forcing_acc_history', {})
val_nav_object_nll_history = infos.get('val_nav_object_nll_history', {})
val_nav_object_teacher_forcing_acc_history = infos.get(
'val_nav_object_teacher_forcing_acc_history', {})
val_nav_room_nll_history = infos.get('val_nav_room_nll_history', {})
val_nav_room_teacher_forcing_acc_history = infos.get(
'val_nav_room_teacher_forcing_acc_history', {})
loss_history = infos.get('loss_history', {})
nll_loss_history = infos.get('nll_loss_history', {})
mse_loss_history = infos.get('mse_loss_history', {})
lr = infos.get('lr', args.learning_rate)
best_val_score, best_val_acc, best_predictions = None, None, None
# start training
while iters <= args.max_iters:
print('Starting epoch %d' % epoch)
# reset seq_length
if args.use_curriculum:
# assume we need 4 epochs to get full seq_length
seq_length = min((args.max_seq_length // 4)**(epoch + 1),
args.max_seq_length)
train_dataset.reset_seq_length(seq_length)
else:
seq_length = args.max_seq_length
# train
for batch in train_loader:
# set mode
model.train()
# zero gradient
optimizer.zero_grad()
# batch = {qid, path_ix, house, id, type, phrase, phrase_emb, ego_feats, next_feats, res_feats,
# action_inputs, action_outputs, action_masks, ego_imgs}
ego_feats = batch['ego_feats'].cuda() # (n, L, 3200)
phrase_embs = batch['phrase_emb'].cuda() # (n, 300)
action_inputs = batch['action_inputs'].cuda() # (n, L)
action_outputs = batch['action_outputs'].cuda() # (n, L)
action_masks = batch['action_masks'].cuda() # (n, L)
# forward
# - logprobs (n, L, #actions)
# - output_feats (n, L, rnn_size)
# - pred_feats (n, L, 3200) or None
logprobs, _, pred_feats, _ = model(ego_feats, phrase_embs,
action_inputs)
nll_loss = nll_crit(logprobs, action_outputs, action_masks)
mse_loss = 0
if args.use_next:
next_feats = batch['next_feats'].cuda() # (n, L, 3200)
mse_loss = mse_crit(pred_feats, next_feats, action_masks)
if args.use_residual:
res_feats = batch['res_feats'].cuda() # (n, L, 3200)
mse_loss = mse_crit(pred_feats, res_feats, action_masks)
loss = nll_loss + args.mse_weight * mse_loss
# backward
loss.backward()
model_utils.clip_gradient(optimizer, args.grad_clip)
optimizer.step()
# training log
if iters % args.losses_log_every == 0:
loss_history[iters] = loss.item()
nll_loss_history[iters] = nll_loss.item()
mse_loss_history[iters] = mse_loss.item() if (
args.use_next or args.use_residual) else 0
print('iters[%s]epoch[%s], train_loss=%.3f (nll_loss=%.3f, mse_loss=%.3f) lr=%.2E, cur_seq_length=%s' % \
(iters, epoch, loss_history[iters], nll_loss_history[iters], mse_loss_history[iters], lr, train_loader.dataset.cur_seq_length))
# decay learning rate
if args.learning_rate_decay_start > 0 and iters > args.learning_rate_decay_start:
frac = (iters - args.learning_rate_decay_start
) / args.learning_rate_decay_every
decay_factor = 0.1**frac
lr = args.learning_rate * decay_factor
model_utils.set_lr(optimizer, lr)
# evaluate
if iters % args.save_checkpoint_every == 0:
print('Checking validation ...')
predictions, overall_nll, overall_teacher_forcing_acc, overall_mse, Nav_nll, Nav_teacher_forcing_acc = \
evaluate(val_dataset, model, nll_crit, mse_crit, opt)
val_nll_history[iters] = overall_nll
val_teacher_forcing_acc_history[
iters] = overall_teacher_forcing_acc
val_mse_history[iters] = overall_mse
val_nav_object_nll_history[iters] = Nav_nll['object']
val_nav_object_teacher_forcing_acc_history[
iters] = Nav_teacher_forcing_acc['object']
val_nav_room_nll_history[iters] = Nav_nll['room']
val_nav_room_teacher_forcing_acc_history[
iters] = Nav_teacher_forcing_acc['room']
# save model if best
# consider all three accuracy, perhaps a better weighting is needed.
current_score = -overall_nll
if best_val_score is None or current_score > best_val_score:
best_val_score = current_score
best_val_acc = overall_teacher_forcing_acc
best_predictions = predictions
checkpoint_path = osp.join(args.checkpoint_dir,
'%s.pth' % args.id)
checkpoint = {}
checkpoint['model_state'] = model.state_dict()
checkpoint['opt'] = vars(args)
torch.save(checkpoint, checkpoint_path)
print('model saved to %s.' % checkpoint_path)
# write to json report
infos['iters'] = iters
infos['epoch'] = epoch
infos['loss_history'] = loss_history
infos['nll_loss_history'] = nll_loss_history
infos['mse_loss_history'] = mse_loss_history
infos['val_nll_history'] = val_nll_history
infos[
'val_teacher_forcing_acc_history'] = val_teacher_forcing_acc_history
infos['val_mse_history'] = val_mse_history
infos[
'val_nav_object_nll_history'] = val_nav_object_nll_history
infos[
'val_nav_object_teacher_forcing_acc_history'] = val_nav_object_teacher_forcing_acc_history
infos['val_nav_room_nll_history'] = val_nav_room_nll_history
infos[
'val_nav_room_teacher_forcing_acc_history'] = val_nav_room_teacher_forcing_acc_history
infos['best_val_score'] = best_val_score
infos['best_val_acc'] = best_val_acc
infos[
'best_predictions'] = predictions if best_predictions is None else best_predictions
infos['opt'] = vars(args)
infos['act_to_ix'] = train_dataset.act_to_ix
infos_json = osp.join(args.checkpoint_dir, '%s.json' % args.id)
with open(infos_json, 'w') as f:
json.dump(infos, f)
print('infos saved to %s.' % infos_json)
# update iters
iters += 1
# update epoch
epoch += 1