def main(args): exp_path = os.path.join(args.output_dir, args.exp_name) log_path = os.path.join(exp_path, 'logs') checkpoint_path = os.path.join(exp_path, 'checkpoints') if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) if os.path.exists(exp_path): print('Error: Experiment already exists, please rename --exp-name') exit() os.makedirs(log_path) os.mkdir(checkpoint_path) print("All experiment outputs will be saved within:", exp_path) # set seed torch.manual_seed(args.seed) np.random.seed(args.seed) # get models and load pre-trained disparity network disp_net = DispNet.DispNet(1).to(device) disp_net.init_weights() if args.disp_net is not None: disp_net.load_state_dict(torch.load(args.disp_net, map_location='cpu')) disp_net.train() pose_net = PoseNet.PoseNet(1, args.sequence_length - 1).to(device) pose_net.init_weights() pose_net.train() # joint optimizer (pose and depth) optim_params = [{ 'params': disp_net.parameters(), 'lr': args.learning_rate }, { 'params': pose_net.parameters(), 'lr': args.learning_rate }] optim = torch.optim.Adam(optim_params, betas=(args.momentum, args.beta), weight_decay=args.weight_decay) # get sequential dataset train_set = CPETDataset.CPET(args.dataset_dir, 'train', args.sequence_length, args.seed) val_set = CPETDataset.CPET(args.dataset_dir, 'val', args.sequence_length, args.seed) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, pin_memory=True) # custom view synthesis loss and depth smoothness loss criterion = ViewSynthesisLoss(device, args.rotation_mode, args.padding_mode) w_synth, w_smooth = args.photo_loss_weight, args.smooth_loss_weight # visualizer visualizer = Visualizer(exp_path, device) # commence experiment print( "Experiment commencing on 4 train seq and 1 validation seq for {} epochs..." .format(args.epochs)) start_time = time.time() # track losses and absolute trajectory error train_loss = np.zeros((args.epochs, 3)) val_loss = np.zeros((args.epochs, 3)) val_ate_mean = np.zeros(args.epochs) total_time = np.zeros(args.epochs) for epo in range(args.epochs): # run training epoch and generate / save random visualizations l_train = train_epoch(disp_net, pose_net, train_loader, criterion, optim, w_synth, w_smooth) train_loss[epo, :] = l_train[:] visualizer.generate_random_visuals(disp_net, pose_net, train_loader, criterion, args.vis_per_epoch, epo, 'train') # run validation epoch and acquire pose estimation metrics. Plot trajectories l_val, ate, ate_mean, gt_traj, pred_traj = validate( disp_net, pose_net, val_loader, criterion, w_synth, w_smooth) val_loss[epo, :] = l_val[:] val_ate_mean[epo] = ate_mean # visualization of disparity maps, BEV trajectories, and 3D trajectories visualizer.generate_random_visuals(disp_net, pose_net, val_loader, criterion, args.vis_per_epoch, epo, 'val') visualizer.generate_trajectory(pred_traj, 'pred', 'Estimated', epo, 'val') visualizer.generate_trajectories(gt_traj, pred_traj, "Horns", epo, 'val') visualizer.generate_3d_trajectory(gt_traj, pred_traj, "Horns", epo, 'val') if epo == 0: visualizer.generate_trajectory(gt_traj, 'gt', 'True', epo, 'val') total_time[epo] = time.time() - start_time print_str = "epo - {}/{} | train_loss - {:.3f} | val_loss - {:.3f} | ".format( epo, args.epochs, train_loss[epo, 0], val_loss[epo, 0]) print_str += "val_ate - {:.3f} | total_time - {}".format( ate_mean, datetime.timedelta(seconds=total_time[epo])) print(print_str) # save models if (epo + 1) % args.save_freq == 0: model_checkpoint(disp_net, 'disp_net_' + str(epo + 1), checkpoint_path) model_checkpoint(pose_net, 'pose_net_' + str(epo + 1), checkpoint_path) # save current stats np.savetxt(os.path.join(log_path, 'train_loss.txt'), train_loss) np.savetxt(os.path.join(log_path, 'val_loss.txt'), val_loss) np.savetxt(os.path.join(log_path, 'val_ate_mean.txt'), val_ate_mean) np.savetxt(os.path.join(log_path, 'time_log.txt'), total_time) # generate metric curves generate_curve([train_loss[:, 0], val_loss[:, 0]], ['train', 'val'], 'loss', 'Train vs Val Combined Loss', log_path) generate_curve([train_loss[:, 1], val_loss[:, 1]], ['train', 'val'], 'photometric loss', 'Train vs Val Photometric Reconstruction Loss', log_path) generate_curve([train_loss[:, 2], val_loss[:, 2]], ['train', 'val'], 'depth smooth loss', 'Train vs Val Depth Smoothness Loss', log_path) generate_curve([val_ate_mean], ['val'], 'ATE', 'Validation Absolute Trajectory Error', log_path)
def main(args): exp_path = os.path.join(args.output_dir, args.exp_name) log_path = os.path.join(exp_path, 'logs') if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) if os.path.exists(exp_path): print('Error: Experiment already exists, over-writing experiment') shutil.rmtree(exp_path) os.makedirs(log_path) print("All experiment outputs will be saved within:", exp_path) # set seed torch.manual_seed(args.seed) np.random.seed(args.seed) # get models, load pre-trained disparity network and pose network disp_net = DispNet.DispNet(1).to(device) disp_net.load_state_dict(torch.load(args.disp_net, map_location='cpu')) disp_net.eval() pose_net = PoseNet.PoseNet(1, args.sequence_length - 1).to(device) pose_net.load_state_dict(torch.load(args.pose_net, map_location='cpu')) pose_net.eval() # get sequence dataset test_set = CPETDataset.CPET(args.dataset_dir, args.run_sequence, args.sequence_length, args.seed) test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.batch_size, shuffle=False, pin_memory=True) # custom view synthesis loss and depth smoothness loss criterion = ViewSynthesisLoss(device, args.rotation_mode, args.padding_mode) w1, w2 = args.photo_loss_weight, args.smooth_loss_weight # visualizer visualizer = Visualizer(exp_path, device) print("Commencing testing on {} sequence...".format(args.run_sequence)) # run test epoch, acquire pose estimation metrics (ATE) from Horn's Method and Umeyama Method start_time = time.time() l_test, horn, umeyama, rate = test(disp_net, pose_net, test_loader, criterion, visualizer, args.skip_freq, w1, w2) total_time = time.time() - start_time # visualize estimated and ground truth trajectories in BEV / 3D - Horn's alignment visualizer.generate_trajectories(horn[2], horn[3], "Horns", epo, args.run_sequence) visualizer.generate_3d_trajectory(horn[2], horn[3], "Horns", epo, args.run_sequence) # visualize estimated and ground truth trajectories in BEV / 3D - Umeyama alignment visualizer.generate_trajectories(umeyama[2], umeyama[3], "Umeyama", epo, args.run_sequence) visualizer.generate_3d_trajectory(umeyama[2], umeyama[3], "Umeyama", epo, args.run_sequence) # visualize trajectories independently - Umeyama visualizer.generate_trajectory(umeyama[2], 'gt', 'True', epo, args.run_sequence) visualizer.generate_trajectory(umeyama[3], 'pred', 'Estimated', epo, args.run_sequence) print_str = "ATE (Umeyama) - {:.3f} | ATE (Horn's) - {:.3f}".format( umeyama[1], horn[1]) print_str += " | view synth loss - {:.3f} | smooth loss - {:.3f}".format( l_test[1], l_test[2]) print_str += " | Hz - {:.5f} | total time - {}".format( rate, datetime.timedelta(seconds=total_time)) print(print_str) # save current stats np.savetxt(os.path.join(log_path, '{}_loss.txt'.format(args.run_sequence)), l_test) np.savetxt( os.path.join(log_path, '{}_ate_mean_horn.txt'.format(args.run_sequence)), np.array([horn[1]])) np.savetxt( os.path.join(log_path, '{}_ate_full_horn.txt'.format(args.run_sequence)), horn[0]) np.savetxt( os.path.join(log_path, '{}_ate_mean_umeyama.txt'.format(args.run_sequence)), np.array([umeyama[1]])) np.savetxt( os.path.join(log_path, '{}_ate_full_umeyama.txt'.format(args.run_sequence)), umeyama[0]) np.savetxt(os.path.join(log_path, 'time_log.txt'), np.array([total_time])) print('-----')