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 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(opt) model.load_state_dict(checkpoint['model_state']) print('model set up.') # evaluate eval(0, args, model, 0, args.split)
def main(args): mp.set_start_method('forkserver', force=True) args.gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',') args.gpus = [int(x) for x in args.gpus] # set up shared_model checkpoint_path = osp.join(args.checkpoint_dir, '%s.pth' % args.start_from) checkpoint = torch.load(checkpoint_path) shared_nav_model = Navigator(checkpoint['opt']) shared_nav_model.load_state_dict(checkpoint['model_state']) shared_nav_model.cpu() shared_nav_model.share_memory() print('shared_nav_model set up.') # some arguments need to be copied from start_from args.use_action = checkpoint['opt']['use_action'] args.nav_types = ['room'] # processes processes = [] counter = mp.Value('i', 0) lock = mp.Lock() # train for rank in range(args.num_processes): p = mp.Process(target=train, args=(rank, args, shared_nav_model, counter, lock)) p.start() processes.append(p) # imitation p = mp.Process(target=imitation, args=(args.num_processes, args, shared_nav_model, counter)) p.start() processes.append(p) # eval p = mp.Process(target=eval, args=(args.num_processes + 1, args, shared_nav_model, counter, 'val')) p.start() processes.append(p) # join for p in processes: p.join()
def eval(rank, args, shared_nav_model, counter, split): # metric_names metric_names = [ 'd_0_5', 'd_0_10', 'd_0_15', 'd_T_5', 'd_T_10', 'd_T_15', 'd_D_5', 'd_D_10', 'd_D_15', 'd_min_5', 'd_min_10', 'd_min_15', 'h_T_5', 'h_T_10', 'h_T_15', 'r_T_5', 'r_T_10', 'r_T_15', 'r_e_5', 'r_e_10', 'r_e_15', 'stop_5', 'stop_10', 'stop_15', 'ep_len_5', 'ep_len_10', 'ep_len_15' ] # set up dataset cfg = { 'colorFile': osp.join(args.house_meta_dir, 'colormap_fine.csv'), 'roomTargetFile': osp.join(args.house_meta_dir, 'room_target_object_map.csv'), 'modelCategoryFile': osp.join(args.house_meta_dir, 'ModelCategoryMapping.csv'), 'prefix': args.house_data_dir, } 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': split, 'max_seq_length': args.max_seq_length, 'nav_types': ['room'], 'gpu_id': 0, 'cfg': cfg, 'max_threads_per_gpu': args.max_threads_per_gpu, 'target_obj_conn_map_dir': args.target_obj_conn_map_dir, 'map_resolution': 500, 'pretrained_cnn_path': args.pretrained_cnn_path, 'requires_imgs': False, 'question_types': ['all'], } dataset = NavReinforceDataset(**loader_kwargs) eval_loader = DataLoader(dataset, batch_size=1, num_workers=0) print('eval_loader set up.') # set up model opt = vars(args) opt['act_to_ix'] = dataset.act_to_ix opt['num_actions'] = len(opt['act_to_ix']) model = Navigator(opt) print('eval navigator set up.') while True: # run evaluation for this epoch invalids = [] # metrics nav_metrics = NavMetric( info={ 'split': split, 'thread': rank }, metric_names=metric_names, log_json=args.results_json, ) # sync model (fixed since now on!) model.load_state_dict(shared_nav_model.state_dict()) model.eval() model.cuda() # reset envs eval_loader.dataset._load_envs(start_idx=0, in_order=True) # run done = False # current epoch is not done yet predictions = [] while not done: for batch in tqdm(eval_loader): # load target_paths from available_idx (of some envs) # batch = {idx, qid, path_ix, house, id, type, phrase, phrase_emb, ego_feats, # action_inputs, action_outputs, action_masks, action_length} idx = batch['idx'][0].item() qid = batch['qid'][0].item() phrase_emb = batch['phrase_emb'].cuda() # (1, 300) float raw_ego_feats = batch['ego_feats'] # (1, L, 3200) float action_inputs = batch['action_inputs'] # (1, L) int action_outputs = batch['action_outputs'] # (1, L) int action_length = batch['action_length'][0].item() tgt_id = batch['id'][0] tgt_type = batch['type'][0] pred = { 'qid': qid, 'house': batch['house'][0], 'id': batch['id'][0], 'type': batch['type'][0], 'path_name': batch['path_name'][0], 'path_ix': batch['path_ix'][0].item(), 'start_ix': batch['start_ix'][0].item(), 'key_ix': batch['key_ix'][0].item(), 'action_length': action_length, 'phrase': batch['phrase'][0], } metrics_slug = {} # evaluate at multiple initializations for i in [5, 10, 15]: if action_length - i < 0: invalids.append((idx, i)) continue h3d = eval_loader.dataset.episode_house episode_length = 0 episode_done = True dists_to_target = [] pos_queue = [] actions = [] # forward through lstm till spawn if len(eval_loader.dataset.episode_pos_queue[:-i]): prev_pos_queue = eval_loader.dataset.episode_pos_queue[: -i] # till spawned position # ego_imgs = eval_loader.dataset.get_frames(h3d, prev_pos_queue, preprocess=True) # ego_imgs = torch.from_numpy(ego_imgs).cuda() # (l, 3, 224, 224) # ego_feats = eval_loader.dataset.cnn.extract_feats(ego_imgs, conv4_only=True) # (l, 3200) # ego_feats = ego_feats.view(1, len(prev_pos_queue), 3200) # (1, l, 3200) ego_feats_pruned = raw_ego_feats[:, :len( prev_pos_queue), :].cuda() # (1, l, 3200) action_inputs_pruned = action_inputs[:, :len( prev_pos_queue)].cuda() # (1, l) _, _, _, state = model( ego_feats_pruned, phrase_emb, action_inputs_pruned) # (1, l, rnn_size) action = action_inputs[0, len(prev_pos_queue)].view( -1).cuda() # (1, ) init_pos = eval_loader.dataset.episode_pos_queue[-i] else: state = None action = torch.LongTensor([ eval_loader.dataset.act_to_ix['dummy'] ]).cuda() # (1, ) init_pos = eval_loader.dataset.episode_pos_queue[-i] # spawn h3d.env.reset(x=init_pos[0], y=init_pos[2], yaw=init_pos[3]) init_dist_to_target = h3d.get_dist_to_target( h3d.env.cam.pos) if init_dist_to_target < 0: # unreachable invalids.append([idx, i]) continue dists_to_target += [init_dist_to_target] pos_queue += [init_pos] # act ego_img = h3d.env.render() ego_img = ( torch.from_numpy(ego_img.transpose(2, 0, 1)).float() / 255.).cuda() ego_feat = eval_loader.dataset.cnn.extract_feats( ego_img.unsqueeze(0), conv4_only=True) # (1, 3200) collision = False rot_act = None rot_act_cnt = 0 for step in range(args.max_episode_length): # forward model episode_length += 1 logprobs, state = model.forward_step( ego_feat, phrase_emb, action, state) # (1, 4), (1, rnn_size) # special case 1: # if previous action is "forward" and collision happend and this action is still "forward", suppress it. if action.item() == 0 and collision and torch.exp( logprobs[0]).argmax().item() == 0: logprobs[0][0] = -1e5 # special case 2: # if spinned around 6 times for same rotation action, we suppress it if torch.exp(logprobs[0]).argmax().item( ) == rot_act and rot_act_cnt > 5: logprobs[0][torch.exp( logprobs[0]).argmax().item()] = -1e5 # sample action action = torch.exp(logprobs[0]).argmax().item() actions += [action] # accumulate rot_act if action == 0: rot_act = None rot_act_cnt = 0 elif action in [1, 2]: if rot_act == action: rot_act_cnt += 1 else: rot_act = action rot_act_cnt = 1 # interact with environment ego_img, _, episode_done, collision = h3d.step(action) # prepare state for next action ego_img = (torch.from_numpy(ego_img.transpose( 2, 0, 1)).float() / 255.).cuda() ego_feat = eval_loader.dataset.cnn.extract_feats( ego_img.unsqueeze(0), conv4_only=True) # (1, 3200) action = torch.LongTensor([action]).cuda() # (1, ) # add to result dists_to_target.append( h3d.get_dist_to_target(h3d.env.cam.pos)) pos_queue.append([ h3d.env.cam.pos.x, h3d.env.cam.pos.y, h3d.env.cam.pos.z, h3d.env.cam.yaw ]) if episode_done: break # add to predictions pred['d_%s' % str(i)] = { 'actions': actions, 'pos_queue': pos_queue, 'gd_actions': action_outputs[ 0, len(prev_pos_queue):len(prev_pos_queue) + i].tolist(), 'gd_pos_queue': eval_loader.dataset.episode_pos_queue[-i:] + [eval_loader.dataset.episode_pos_queue[-1]] } # is final view looking at target object? if tgt_type == 'object': iou = h3d.compute_target_iou(tgt_id) R = 1 if h3d.compute_target_iou(tgt_id) >= 0.1 else 0 pred['iou_%s' % str(i)] = iou # compute stats metrics_slug['d_0_' + str(i)] = dists_to_target[0] metrics_slug['ep_len_' + str(i)] = episode_length metrics_slug['stop_' + str(i)] = 1 if action == 3 else 0 if tgt_type == 'object': metrics_slug['d_T_' + str(i)] = dists_to_target[-1] metrics_slug[ 'd_D_' + str(i)] = dists_to_target[0] - dists_to_target[-1] metrics_slug['d_min_' + str(i)] = float( np.array(dists_to_target).min()) metrics_slug['h_T_' + str(i)] = R else: inside_room = [] for p in pos_queue: inside_room.append( h3d.is_inside_room( p, eval_loader.dataset.target_room)) if inside_room[-1] == True: metrics_slug['r_T_' + str(i)] = 1 else: metrics_slug['r_T_' + str(i)] = 0 if any([x == True for x in inside_room]) == True: metrics_slug['r_e_' + str(i)] = 1 else: metrics_slug['r_e_' + str(i)] = 0 # collate and update metrics metrics_list = [] for i in nav_metrics.metric_names: if i not in metrics_slug: metrics_list.append(nav_metrics.metrics[ nav_metrics.metric_names.index(i)][0]) else: metrics_list.append(metrics_slug[i]) # update metrics if len(metrics_slug) > 0: nav_metrics.update(metrics_list) predictions.append(pred) print(nav_metrics.get_stat_string(mode=0)) print('invalids', len(invalids)) logging.info("EVAL: init_steps: {} metrics: {}".format( i, nav_metrics.get_stat_string(mode=0))) logging.info("EVAL: init_steps: {} invalids: {}".format( i, len(invalids))) # next environments eval_loader.dataset._load_envs() print("eval_loader pruned_env_set len: {}".format( len(eval_loader.dataset.pruned_env_set))) logging.info("eval_loader pruned_env_set len: {}".format( len(eval_loader.dataset.pruned_env_set))) if len(eval_loader.dataset.pruned_env_set) == 0: done = True # save results nav_metrics.dump_log(predictions) # write summary results_str = 'id[%s] split[%s]\n' % (args.id, args.split) for n, r in zip(metric_names, nav_metrics.stats[-1]): if r: results_str += '%10s: %6.2f,' % (n, r[0]) if '_15' in n: results_str += '\n' f = open(args.report_txt, 'a') f.write(results_str) f.write('\n') f.close() # break break
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 eval(rank, args, shared_nav_model, counter, split='val'): # set up cuda device torch.cuda.set_device(args.gpus.index(args.gpus[rank % len(args.gpus)])) # tensorboard # writer = tb.SummaryWriter(log_dir=args.tb_dir, filename_suffix=str(rank)) writer = tb.SummaryWriter(osp.join(args.tb_dir, str(rank))) # set up dataset cfg = { 'colorFile': osp.join(args.house_meta_dir, 'colormap_fine.csv'), 'roomTargetFile': osp.join(args.house_meta_dir, 'room_target_object_map.csv'), 'modelCategoryFile': osp.join(args.house_meta_dir, 'ModelCategoryMapping.csv'), 'prefix': args.house_data_dir, } 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': split, 'max_seq_length': args.max_seq_length, 'nav_types': args.nav_types, 'gpu_id': 0, 'cfg': cfg, 'max_threads_per_gpu': args.max_threads_per_gpu, 'target_obj_conn_map_dir': args.target_obj_conn_map_dir, 'map_resolution': 500, 'pretrained_cnn_path': args.pretrained_cnn_path, 'requires_imgs': False, 'question_types': ['all'], } dataset = NavReinforceDataset(**loader_kwargs) eval_loader = DataLoader(dataset, batch_size=1, num_workers=0) print('eval_loader set up.') # set up model opt = vars(args) opt['act_to_ix'] = dataset.act_to_ix opt['num_actions'] = len(opt['act_to_ix']) model = Navigator(opt) print('eval navigator set up.') # evaluate epoch = 0 best_eval_score = None while epoch < int(args.max_epochs): # wait till counter.value >= epoch * num_iters_per_epoch cur_counter_value = counter.value if counter.value / args.num_iters_per_epoch >= epoch: epoch += 1 else: continue # run evaluation for this epoch invalids = [] # metrics nav_metrics = NavMetric( info={ 'split': split, 'thread': rank }, metric_names=metric_names, log_json=osp.join(args.log_dir, 'nav_eval_val.json'), ) # update model (fixed since now on!) model.load_state_dict(shared_nav_model.state_dict()) model.eval() model.cuda() # reset envs eval_loader.dataset._load_envs(start_idx=0, in_order=True) eval_loader.dataset.visited_envs = set() # run done = False predictions = [] while not done: for batch in tqdm(eval_loader): # load target_paths from available_idx (of some envs) # batch = {idx, qid, path_ix, house, id, type, phrase, phrase_emb, ego_feats, # action_inputs, action_outputs, action_masks, action_length} idx = batch['idx'][0].item() qid = batch['qid'][0].item() phrase_emb = batch['phrase_emb'].cuda() # (1, 300) cuda raw_ego_feats = batch['ego_feats'] # (1, L, 3200) float action_inputs = batch['action_inputs'] # (1, L) int action_outputs = batch['action_outputs'] # (1, L) int action_length = batch['action_length'][0].item() tgt_type = batch['type'][0] tgt_id = batch['id'][0] pred = { 'qid': qid, 'house': batch['house'][0], 'id': batch['id'][0], 'type': batch['type'][0], 'path_name': batch['path_name'][0], 'path_ix': batch['path_ix'][0].item(), 'start_ix': batch['start_ix'][0].item(), 'key_ix': batch['key_ix'][0].item(), 'action_length': action_length, 'phrase': batch['phrase'][0], } metrics_slug = {} # evaluate at multiple initializations # for i in [5, 10, 15]: for i in [15]: # for saving evaluation time if action_length - i < 0: invalids.append((idx, i)) continue h3d = eval_loader.dataset.episode_house episode_length = 0 episode_done = True dists_to_target = [] pos_queue = [] actions = [] # forward through navigator till spawn if len(eval_loader.dataset.episode_pos_queue[:-i]): prev_pos_queue = eval_loader.dataset.episode_pos_queue[: -i] # till spawned position # ego_imgs = eval_loader.dataset.get_frames(h3d, prev_pos_queue, preprocess=True) # ego_imgs = torch.from_numpy(ego_imgs).cuda() # (l, 3, 224, 224) # ego_feats = eval_loader.dataset.cnn.extract_feats(ego_imgs, conv4_only=True) # (l, 3200) # ego_feats = ego_feats.view(1, len(prev_pos_queue), 3200) # (1, l, 3200) ego_feats_pruned = raw_ego_feats[:, :len( prev_pos_queue), :].cuda() # (1, l, 3200) action_inputs_pruned = action_inputs[:, :len( prev_pos_queue)].cuda() # (1, l) _, _, _, state = model( ego_feats_pruned, phrase_emb, action_inputs_pruned) # (1, l, rnn_size) action = action_inputs[0, len(prev_pos_queue)].view( -1).cuda() # (1, ) init_pos = eval_loader.dataset.episode_pos_queue[-i] else: state = None action = torch.LongTensor([ eval_loader.dataset.act_to_ix['dummy'] ]).cuda() # (1, ) init_pos = eval_loader.dataset.episode_pos_queue[ 0] # use first position instead # spawn h3d.env.reset(x=init_pos[0], y=init_pos[2], yaw=init_pos[3]) init_dist_to_target = h3d.get_dist_to_target( h3d.env.cam.pos) if init_dist_to_target < 0: # unreachable invalids.append([idx, i]) continue dists_to_target += [init_dist_to_target] pos_queue += [init_pos] # act ego_img = h3d.env.render() ego_img = ( torch.from_numpy(ego_img.transpose(2, 0, 1)).float() / 255.).cuda() ego_feat = eval_loader.dataset.cnn.extract_feats( ego_img.unsqueeze(0), conv4_only=True) # (1, 3200) collision = False rot_act = None rot_act_cnt = 0 for step in range(args.max_episode_length): # forward model episode_length += 1 with torch.no_grad(): logprobs, state = model.forward_step( ego_feat, phrase_emb, action, state) # (1, 4), (1, rnn_size) # special case 1: if previous action is "forward" and collision happend and this action # is still "forward", suppress it. if action.item() == 0 and collision and torch.exp( logprobs[0]).argmax().item() == 0: logprobs[0][0] = -1e5 # special case 2: # if spinned around 6 times for same rotation action, we suppress it if torch.exp(logprobs[0]).argmax().item( ) == rot_act and rot_act_cnt > 5: logprobs[0][torch.exp( logprobs[0]).argmax().item()] = -1e5 # sample action action = torch.exp(logprobs[0]).argmax().item() actions += [action] # accumulate rot_act if action == 0: rot_act = None rot_act_cnt = 0 elif action in [1, 2]: if rot_act == action: rot_act_cnt += 1 else: rot_act = action rot_act_cnt = 1 # interact with environment ego_img, _, episode_done, collision = h3d.step(action) episode_done = episode_done or episode_length >= args.max_episode_length # no need actually # prepare state for next action ego_img = (torch.from_numpy(ego_img.transpose( 2, 0, 1)).float() / 255.).cuda() ego_feat = eval_loader.dataset.cnn.extract_feats( ego_img.unsqueeze(0), conv4_only=True) # (1, 3200) action = torch.LongTensor([action]).cuda() # (1, ) # add to result dists_to_target.append( h3d.get_dist_to_target(h3d.env.cam.pos)) pos_queue.append([ h3d.env.cam.pos.x, h3d.env.cam.pos.y, h3d.env.cam.pos.z, h3d.env.cam.yaw ]) if episode_done: break # add to predictions pred['d_%s' % str(i)] = { 'actions': actions, 'pos_queue': pos_queue, 'gd_actions': action_outputs[ 0, len(prev_pos_queue):len(prev_pos_queue) + i].tolist(), 'gd_pos_queue': eval_loader.dataset.episode_pos_queue[-i:] + [eval_loader.dataset.episode_pos_queue[-1]] } # compute stats metrics_slug['d_0_' + str(i)] = dists_to_target[0] metrics_slug['ep_len_' + str(i)] = episode_length metrics_slug['stop_' + str(i)] = 1 if action == 3 else 0 if tgt_type == 'object': metrics_slug['d_T_' + str(i)] = dists_to_target[-1] metrics_slug[ 'd_D_' + str(i)] = dists_to_target[0] - dists_to_target[-1] metrics_slug['d_min_' + str(i)] = float( np.array(dists_to_target).min()) iou = h3d.compute_target_iou(tgt_id) metrics_slug['h_T_' + str(i)] = 1 if iou >= 0.1 else 0 else: inside_room = [] for p in pos_queue: inside_room.append( h3d.is_inside_room( p, eval_loader.dataset.target_room)) metrics_slug['r_T_' + str(i)] = 1 if inside_room[-1] else 0 metrics_slug['r_e_' + str(i)] = 1 if any( [x == True for x in inside_room]) else 0 # collate and update metrics metrics_list = [] for name in nav_metrics.metric_names: if name not in metrics_slug: metrics_list.append(nav_metrics.metrics[ nav_metrics.metric_names.index(name)][0]) else: metrics_list.append(metrics_slug[name]) # update metrics if len(metrics_slug) > 0: nav_metrics.update(metrics_list) predictions.append(pred) print(nav_metrics.get_stat_string(mode=0)) print('invalids', len(invalids)) logging.info("EVAL: init_steps: {} metrics: {}".format( i, nav_metrics.get_stat_string(mode=0))) logging.info("EVAL: init_steps: {} invalids: {}".format( i, len(invalids))) print('%s/%s envs visited.' % (len(eval_loader.dataset.visited_envs), len(eval_loader.dataset.env_set))) # next environments eval_loader.dataset._load_envs() print("eval_loader pruned_env_set len: {}".format( len(eval_loader.dataset.pruned_env_set))) logging.info("eval_loader pruned_env_set len: {}".format( len(eval_loader.dataset.pruned_env_set))) if len(eval_loader.dataset.pruned_env_set) == 0: done = True # write to tensorboard for metric_name in [ 'd_T_5', 'd_T_10', 'd_T_15', 'd_D_5', 'd_D_10', 'd_D_15', 'h_T_5', 'h_T_10', 'h_T_15', 'r_T_5', 'r_T_10', 'r_T_15', 'r_e_5', 'r_e_10', 'r_e_15' ]: value = nav_metrics.metrics[nav_metrics.metric_names.index( metric_name)][0] if value: # instead of counter.value (as we started eval at cur_counter_value) writer.add_scalar('eval/%s' % metric_name, value, cur_counter_value) # save if best # best_score = d_D_15 + h_T_15 + r_T_15 cur_score = 0 if nav_metrics.metrics[nav_metrics.metric_names.index('d_D_15')][0]: cur_score += nav_metrics.metrics[nav_metrics.metric_names.index( 'd_D_15')][0] if nav_metrics.metrics[nav_metrics.metric_names.index('h_T_15')][0]: cur_score += nav_metrics.metrics[nav_metrics.metric_names.index( 'h_T_15')][0] if nav_metrics.metrics[nav_metrics.metric_names.index('r_T_15')][0]: cur_score += nav_metrics.metrics[nav_metrics.metric_names.index( 'r_T_15')][0] if (best_eval_score is None or cur_score > best_eval_score) and cur_counter_value > 50000: best_eval_score = cur_score nav_metrics.dump_log(predictions) checkpoint_path = osp.join(args.checkpoint_dir, '%s.pth' % args.id) checkpoint = {} checkpoint['model_state'] = model.cpu().state_dict() checkpoint['opt'] = vars(args) torch.save(checkpoint, checkpoint_path) print('model saved to %s.' % checkpoint_path)
def train(rank, args, shared_nav_model, counter, lock): # set up tensorboard # writer = tb.SummaryWriter(args.tb_dir, filename_suffix=str(rank)) writer = tb.SummaryWriter(osp.join(args.tb_dir, str(rank))) # set up cuda device gpu_id = args.gpus.index(args.gpus[rank % len(args.gpus)]) torch.cuda.set_device(gpu_id) # set up random seeds random.seed(args.seed + rank) np.random.randn(args.seed + rank) torch.manual_seed(args.seed + rank) # set up dataset cfg = { 'colorFile': osp.join(args.house_meta_dir, 'colormap_fine.csv'), 'roomTargetFile': osp.join(args.house_meta_dir, 'room_target_object_map.csv'), 'modelCategoryFile': osp.join(args.house_meta_dir, 'ModelCategoryMapping.csv'), 'prefix': args.house_data_dir, } 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, 'nav_types': args.nav_types, 'gpu_id': args.gpus[rank % len(args.gpus)], 'cfg': cfg, 'max_threads_per_gpu': args.max_threads_per_gpu, 'target_obj_conn_map_dir': args.target_obj_conn_map_dir, 'map_resolution': 500, 'pretrained_cnn_path': args.pretrained_cnn_path, 'requires_imgs': False, 'question_types': ['all'], 'ratio': [rank / args.num_processes, (rank + 1) / args.num_processes] } dataset = NavReinforceDataset(**loader_kwargs) train_loader = DataLoader(dataset, batch_size=1, num_workers=0) print('train_loader set up.') # set up optimizer on shared_nav_model 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 model opt = vars(args) opt['act_to_ix'] = dataset.act_to_ix opt['num_actions'] = len(opt['act_to_ix']) model = Navigator(opt) print('navigator[%s] set up.' % rank) # set up metrics outside epoch, as we use running rewards through whole training nav_metrics = NavMetric( info={ 'split': 'train', 'thread': rank }, metric_names=['reward', 'episode_length'], log_json=osp.join(args.log_dir, 'nav_train_' + str(rank) + '.json'), ) nav_metrics.update([0, 100]) reward_list, episode_length_list = [], [] rwd_stats = Statistics() # computing running mean and std of rewards # path length multiplier min_dist = 5 if args.nav_types == [ 'object' ] else 15 # 15 for rl3 and rl5, 10 for rl4 # max_dist = 25 if args.nav_types == ['object'] else 40 max_dist = 35 if args.nav_types == ['object'] else 50 mult = 0.1 if args.nav_types == ['object'] else 0.15 rwd_thresh = 0.1 if args.nav_types == ['object'] else 0.00 epoch = 0 iters = 0 # train while True: # reset envs train_loader.dataset._load_envs(start_idx=0, in_order=True) done = False # current epoch is not done yet while not done: for batch in train_loader: # sync model model.load_state_dict(shared_nav_model.state_dict()) model.train() model.cuda() # load target_paths from available_idx (of some envs) # batch = {idx, qid, path_ix, house, id, type, phrase, phrase_emb, ego_feats, # action_inputs, action_outputs, action_masks, action_length} idx = batch['idx'][0].item() qid = batch['qid'][0].item() phrase_emb = batch['phrase_emb'].cuda() # (1, 300) float raw_ego_feats = batch['ego_feats'] # (1, L, 3200) float action_inputs = batch['action_inputs'] # (1, L) int action_outputs = batch['action_outputs'] # (1, L) int action_length = batch['action_length'][0].item() tgt_type = batch['type'][0] tgt_id = batch['id'][0] tgt_phrase = batch['phrase'][0] # to be recorded episode_length = 0 episode_done = True dists_to_target = [] pos_queue = [] actions = [] rewards = [] nav_log_probs = [] # spawn agent h3d = train_loader.dataset.episode_house if np.random.uniform(0, 1, 1)[0] <= args.shortest_path_ratio: # half chance we use shortest path to spawn the agent (if vlen > 0, i.e., shortest path long enough) use_shortest_path = True vlen = min(max(min_dist, int(mult * action_length)), action_length) # forward throught navigator till spawn if len(train_loader.dataset.episode_pos_queue) > vlen: prev_pos_queue = train_loader.dataset.episode_pos_queue[: -vlen] # till spawned position ego_feats_pruned = raw_ego_feats[:, :len( prev_pos_queue), :].cuda() # (1, l, 3200) action_inputs_pruned = action_inputs[:, :len( prev_pos_queue)].cuda() # (1, l) _, _, _, state = model( ego_feats_pruned, phrase_emb, action_inputs_pruned) # (1, l, rnn_size) action = action_inputs[0, len(prev_pos_queue)].view( -1).cuda() # (1, ) init_pos = train_loader.dataset.episode_pos_queue[ -vlen] else: state = None action = torch.LongTensor([ train_loader.dataset.act_to_ix['dummy'] ]).cuda() # (1, ) init_pos = train_loader.dataset.episode_pos_queue[ 0] # use first position of the path else: # half chance we randomly spawn agent use_shortest_path = False state = None action = torch.LongTensor([ train_loader.dataset.act_to_ix['dummy'] ]).cuda() # (1, ) init_pos, vlen = train_loader.dataset.spawn_agent( min_dist, max_dist) if init_pos is None: # init_pos not found continue # initiate h3d.env.reset(x=init_pos[0], y=init_pos[2], yaw=init_pos[3]) init_dist_to_target = h3d.get_dist_to_target(h3d.env.cam.pos) if init_dist_to_target < 0: # unreachable continue dists_to_target += [init_dist_to_target] pos_queue += [init_pos] # act ego_img = h3d.env.render() ego_img = ( torch.from_numpy(ego_img.transpose(2, 0, 1)).float() / 255.).cuda() ego_feat = train_loader.dataset.cnn.extract_feats( ego_img.unsqueeze(0), conv4_only=True) # (1, 3200) prev_action, collision = None, False for step in range(args.max_episode_length): # forward model one step episode_length += 1 logprobs, state = model.forward_step( ego_feat, phrase_emb, action, state) # (1, 4), (1, rnn_size) # sample action probs = torch.exp(logprobs) # (1, 4) action = probs.multinomial( num_samples=1).detach() # (1, 1) if prev_action == 0 and collision and action[0][0].item( ) == 0: # special case: prev_action == "forward" && collision && cur_action == "forward" # we sample from {'left', 'right', 'stop'} only action = probs[0:1, 1:].multinomial( num_samples=1).detach() + 1 # (1, 1) if len(pos_queue) < min_dist: # special case: our room navigator tends to stop early, let's push it to explore longer action = probs[0:1, :3].multinomial( num_samples=1).detach() # (1, 1) nav_log_probs.append(logprobs.gather(1, action)) # (1, 1) action = action.view(-1) # (1, ) actions.append(action) # interact with environment ego_img, reward, episode_done, collision = h3d.step( action[0].item(), step_reward=True) if not episode_done: reward -= 0.01 # we don't wanna too long trajectory episode_done = episode_done or episode_length >= args.max_episode_length # no need actually reward = max(min(reward, 1), -1) rewards.append(reward) prev_action = action[0].item() # prepare state for next action ego_img = ( torch.from_numpy(ego_img.transpose(2, 0, 1)).float() / 255.).cuda() ego_feat = train_loader.dataset.cnn.extract_feats( ego_img.unsqueeze(0), conv4_only=True) # (1, 3200) # add to result dists_to_target.append( h3d.get_dist_to_target(h3d.env.cam.pos)) pos_queue.append([ h3d.env.cam.pos.x, h3d.env.cam.pos.y, h3d.env.cam.pos.z, h3d.env.cam.yaw ]) if episode_done: break # final reward R = 0 if tgt_type == 'object': R = 1.0 if h3d.compute_target_iou(tgt_id) >= 0.1 else -1.0 else: R = 0.2 if h3d.is_inside_room( pos_queue[-1], train_loader.dataset.target_room) else -0.2 if R > 0 and h3d.compute_room_targets_iou( room_to_objects[tgt_phrase]) > 0.1: R += 1.0 # encourage agent to move to room-specific targets # backward nav_loss = 0 new_rewards = [ ] # recording reshaped rewards, to be computed as moving average advantages = [] for i in reversed(range(len(rewards))): R = 0.99 * R + rewards[i] new_rewards.insert(0, R) rwd_stats.push(R) for nav_log_prob, R in zip(nav_log_probs, new_rewards): advantage = (R - rwd_stats.mean()) / ( rwd_stats.stddev() + 1e-5) # rl2, rl3 # advantage = R - rwd_stats.mean() # rl1 nav_loss = nav_loss - nav_log_prob * advantage advantages.insert(0, advantage) nav_loss /= max(1, len(nav_log_probs)) optimizer.zero_grad() nav_loss.backward() clip_model_gradient(model.parameters(), args.grad_clip) # Till this point, we have grads on model's parameters! # but shared_nav_model's grads is still not updated ensure_shared_grads(model.cpu(), shared_nav_model) optimizer.step() if iters % 5 == 0: log_info = 'train-r%2s(ep%2s it%5s lr%.2E mult%.2f, run_rwd%6.3f, bs_rwd%6.3f, bs_std%2.2f), vlen(%s)=%2s, rwd=%6.3f, ep_len=%3s, tgt:%s' % \ (rank, epoch, iters, lr, mult, nav_metrics.metrics[0][1], rwd_stats.mean(), rwd_stats.stddev(), 'short' if use_shortest_path else 'spawn', vlen, np.mean(new_rewards), len(new_rewards), tgt_phrase) print(log_info) # update metrics reward_list += new_rewards episode_length_list.append(episode_length) if len(episode_length_list) > 50: nav_metrics.update([reward_list, episode_length_list]) nav_metrics.dump_log() reward_list, episode_length_list = [], [] # write to tensorboard writer.add_scalar('train_rank%s/nav_loss' % rank, nav_loss.item(), counter.value) writer.add_scalar('train_rank%s/steps_avg_rwd' % rank, float(np.mean(new_rewards)), counter.value) writer.add_scalar('train_rank%s/steps_sum_rwd' % rank, float(np.sum(new_rewards)), counter.value) writer.add_scalar('train_rank%s/advantage' % rank, float(np.mean(advantages)), counter.value) # increase counter as this episode ends with lock: counter.value += 1 # increase mult iters += 1 if nav_metrics.metrics[0][ 1] > rwd_thresh: # baseline = nav_metrics.metrics[0][1] mult = min(mult + 0.1, 1.0) rwd_thresh += 0.01 else: mult = max(mult - 0.1, 0.1) rwd_thresh -= 0.01 rwd_thresh = max(0.1, min(rwd_thresh, 0.2)) # 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) # next environments train_loader.dataset._load_envs(in_order=True) print("train_loader pruned_env_set len: {}".format( len(train_loader.dataset.pruned_env_set))) logging.info("train_loader pruned_env_set len: {}".format( len(train_loader.dataset.pruned_env_set))) if len(train_loader.dataset.pruned_env_set) == 0: done = True epoch += 1