loss, optimizer, dataloader, pair_generation_tnf, log_interval=100) model.eval() test_loss[epoch - 1] = process_epoch('test', epoch, model, loss, optimizer, dataloader_test, pair_generation_tnf, log_interval=100) # remember best loss is_best = test_loss[epoch - 1] < best_test_loss best_test_loss = min(test_loss[epoch - 1], best_test_loss) save_checkpoint( { 'epoch': epoch + 1, 'args': args, 'state_dict': model.state_dict(), 'best_test_loss': best_test_loss, 'optimizer': optimizer.state_dict(), 'train_loss': train_loss, 'test_loss': test_loss, }, is_best, checkpoint_name) print('Done!')
def main(): args, arg_groups = ArgumentParser(mode='train').parse() print(args) use_cuda = torch.cuda.is_available() device = torch.device('cuda') if use_cuda else torch.device('cpu') # Seed torch.manual_seed(args.seed) if use_cuda: torch.cuda.manual_seed(args.seed) # Download dataset if needed and set paths if args.training_dataset == 'pascal': if args.dataset_image_path == '' and not os.path.exists( 'datasets/pascal-voc11/TrainVal'): download_pascal('datasets/pascal-voc11/') if args.dataset_image_path == '': args.dataset_image_path = 'datasets/pascal-voc11/' args.dataset_csv_path = 'training_data/pascal-random' # CNN model and loss print('Creating CNN model...') if args.geometric_model == 'affine': cnn_output_dim = 6 elif args.geometric_model == 'hom' and args.four_point_hom: cnn_output_dim = 8 elif args.geometric_model == 'hom' and not args.four_point_hom: cnn_output_dim = 9 elif args.geometric_model == 'tps': cnn_output_dim = 18 model = CNNGeometric(use_cuda=use_cuda, output_dim=cnn_output_dim, **arg_groups['model']) if args.geometric_model == 'hom' and not args.four_point_hom: init_theta = torch.tensor([1, 0, 0, 0, 1, 0, 0, 0, 1], device=device) model.FeatureRegression.linear.bias.data += init_theta if args.geometric_model == 'hom' and args.four_point_hom: init_theta = torch.tensor([-1, -1, 1, 1, -1, 1, -1, 1], device=device) model.FeatureRegression.linear.bias.data += init_theta if args.use_mse_loss: print('Using MSE loss...') loss = nn.MSELoss() else: print('Using grid loss...') loss = TransformedGridLoss(use_cuda=use_cuda, geometric_model=args.geometric_model) # Initialize Dataset objects dataset = SynthDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='train.csv', dataset_image_path=args.dataset_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) dataset_val = SynthDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='val.csv', dataset_image_path=args.dataset_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) # Set Tnf pair generation func pair_generation_tnf = SynthPairTnf(geometric_model=args.geometric_model, use_cuda=use_cuda) # Initialize DataLoaders dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4) dataloader_val = DataLoader(dataset_val, batch_size=args.batch_size, shuffle=True, num_workers=4) # Optimizer and eventual scheduler optimizer = optim.Adam(model.FeatureRegression.parameters(), lr=args.lr) if args.lr_scheduler: scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.lr_max_iter, eta_min=1e-6) # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min') else: scheduler = False # Train # Set up names for checkpoints if args.use_mse_loss: ckpt = args.trained_model_fn + '_' + args.geometric_model + '_mse_loss' + args.feature_extraction_cnn checkpoint_path = os.path.join(args.trained_model_dir, args.trained_model_fn, ckpt + '.pth.tar') else: ckpt = args.trained_model_fn + '_' + args.geometric_model + '_grid_loss' + args.feature_extraction_cnn checkpoint_path = os.path.join(args.trained_model_dir, args.trained_model_fn, ckpt + '.pth.tar') if not os.path.exists(args.trained_model_dir): os.mkdir(args.trained_model_dir) # Set up TensorBoard writer if not args.log_dir: tb_dir = os.path.join(args.trained_model_dir, args.trained_model_fn + '_tb_logs') else: tb_dir = os.path.join(args.log_dir, args.trained_model_fn + '_tb_logs') logs_writer = SummaryWriter(tb_dir) # add graph, to do so we have to generate a dummy input to pass along with the graph dummy_input = { 'source_image': torch.rand([args.batch_size, 3, 240, 240], device=device), 'target_image': torch.rand([args.batch_size, 3, 240, 240], device=device), 'theta_GT': torch.rand([16, 2, 3], device=device) } logs_writer.add_graph(model, dummy_input) # Start of training print('Starting training...') best_val_loss = float("inf") for epoch in range(1, args.num_epochs + 1): # we don't need the average epoch loss so we assign it to _ _ = train(epoch, model, loss, optimizer, dataloader, pair_generation_tnf, log_interval=args.log_interval, scheduler=scheduler, tb_writer=logs_writer) val_loss = validate_model(model, loss, dataloader_val, pair_generation_tnf, epoch, logs_writer) # remember best loss is_best = val_loss < best_val_loss best_val_loss = min(val_loss, best_val_loss) save_checkpoint( { 'epoch': epoch + 1, 'args': args, 'state_dict': model.state_dict(), 'best_val_loss': best_val_loss, 'optimizer': optimizer.state_dict(), }, is_best, checkpoint_path) logs_writer.close() print('Done!')
print(mode.capitalize()+' set: Average loss: {:.4f}'.format(epoch_loss)) return epoch_loss train_loss = np.zeros(args.num_epochs) test_loss = np.zeros(args.num_epochs) print('Starting training...') model.FeatureExtraction.eval() for epoch in range(1, args.num_epochs+1): model.FeatureRegression.train() train_loss[epoch-1] = process_epoch('train',epoch,model,loss,optimizer,dataloader,pair_generation_tnf,log_interval=100) model.FeatureRegression.eval() test_loss[epoch-1] = process_epoch('test',epoch,model,loss,optimizer,dataloader_test,pair_generation_tnf,log_interval=100) # remember best loss is_best = test_loss[epoch-1] < best_test_loss best_test_loss = min(test_loss[epoch-1], best_test_loss) save_checkpoint({ 'epoch': epoch + 1, 'args': args, 'state_dict': model.state_dict(), 'best_test_loss': best_test_loss, 'optimizer' : optimizer.state_dict(), 'train_loss': train_loss, 'test_loss': test_loss, }, is_best,checkpoint_name) print('Done!')
def main(): args, arg_groups = ArgumentParser(mode='train').parse() print(args) use_cuda = torch.cuda.is_available() use_me = args.use_me device = torch.device('cuda') if use_cuda else torch.device('cpu') # Seed # torch.manual_seed(args.seed) # if use_cuda: # torch.cuda.manual_seed(args.seed) # CNN model and loss print('Creating CNN model...') if args.geometric_model == 'affine_simple': cnn_output_dim = 3 elif args.geometric_model == 'affine_simple_4': cnn_output_dim = 4 else: raise NotImplementedError('Specified geometric model is unsupported') model = CNNGeometric(use_cuda=use_cuda, output_dim=cnn_output_dim, **arg_groups['model']) if args.geometric_model == 'affine_simple': init_theta = torch.tensor([0.0, 1.0, 0.0], device=device) elif args.geometric_model == 'affine_simple_4': init_theta = torch.tensor([0.0, 1.0, 0.0, 0.0], device=device) try: model.FeatureRegression.linear.bias.data += init_theta except: model.FeatureRegression.resnet.fc.bias.data += init_theta args.load_images = False if args.loss == 'mse': print('Using MSE loss...') loss = nn.MSELoss() elif args.loss == 'weighted_mse': print('Using weighted MSE loss...') loss = WeightedMSELoss(use_cuda=use_cuda) elif args.loss == 'reconstruction': print('Using reconstruction loss...') loss = ReconstructionLoss( int(np.rint(args.input_width * (1 - args.crop_factor) / 16) * 16), int(np.rint(args.input_height * (1 - args.crop_factor) / 16) * 16), args.input_height, use_cuda=use_cuda) args.load_images = True elif args.loss == 'combined': print('Using combined loss...') loss = CombinedLoss(args, use_cuda=use_cuda) if args.use_reconstruction_loss: args.load_images = True elif args.loss == 'grid': print('Using grid loss...') loss = SequentialGridLoss(use_cuda=use_cuda) else: raise NotImplementedError('Specifyed loss %s is not supported' % args.loss) # Initialize Dataset objects if use_me: dataset = MEDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='train.csv', dataset_image_path=args.dataset_image_path, input_height=args.input_height, input_width=args.input_width, crop=args.crop_factor, use_conf=args.use_conf, use_random_patch=args.use_random_patch, normalize_inputs=args.normalize_inputs, random_sample=args.random_sample, load_images=args.load_images) dataset_val = MEDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='val.csv', dataset_image_path=args.dataset_image_path, input_height=args.input_height, input_width=args.input_width, crop=args.crop_factor, use_conf=args.use_conf, use_random_patch=args.use_random_patch, normalize_inputs=args.normalize_inputs, random_sample=args.random_sample, load_images=args.load_images) else: dataset = SynthDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='train.csv', dataset_image_path=args.dataset_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) dataset_val = SynthDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='val.csv', dataset_image_path=args.dataset_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) # Set Tnf pair generation func if use_me: pair_generation_tnf = BatchTensorToVars(use_cuda=use_cuda) elif args.geometric_model == 'affine_simple' or args.geometric_model == 'affine_simple_4': pair_generation_tnf = SynthPairTnf(geometric_model='affine', use_cuda=use_cuda) else: raise NotImplementedError('Specified geometric model is unsupported') # Initialize DataLoaders dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4) dataloader_val = DataLoader(dataset_val, batch_size=args.batch_size, shuffle=True, num_workers=4) # Optimizer optimizer = optim.Adam(model.FeatureRegression.parameters(), lr=args.lr) # Train # Set up names for checkpoints ckpt = args.trained_model_fn + '_' + args.geometric_model + '_' + args.loss + '_loss_' checkpoint_path = os.path.join(args.trained_model_dir, args.trained_model_fn, ckpt + '.pth.tar') if not os.path.exists(args.trained_model_dir): os.mkdir(args.trained_model_dir) # Set up TensorBoard writer if not args.log_dir: tb_dir = os.path.join(args.trained_model_dir, args.trained_model_fn + '_tb_logs') else: tb_dir = os.path.join(args.log_dir, args.trained_model_fn + '_tb_logs') logs_writer = SummaryWriter(tb_dir) # add graph, to do so we have to generate a dummy input to pass along with the graph if use_me: dummy_input = { 'mv_L2R': torch.rand([args.batch_size, 2, 216, 384], device=device), 'mv_R2L': torch.rand([args.batch_size, 2, 216, 384], device=device), 'grid_L2R': torch.rand([args.batch_size, 2, 216, 384], device=device), 'grid_R2L': torch.rand([args.batch_size, 2, 216, 384], device=device), 'grid': torch.rand([args.batch_size, 2, 216, 384], device=device), 'conf_L': torch.rand([args.batch_size, 1, 216, 384], device=device), 'conf_R': torch.rand([args.batch_size, 1, 216, 384], device=device), 'theta_GT': torch.rand([args.batch_size, 4], device=device), } if args.load_images: dummy_input['img_R_orig'] = torch.rand( [args.batch_size, 1, 216, 384], device=device) dummy_input['img_R'] = torch.rand([args.batch_size, 1, 216, 384], device=device) else: dummy_input = { 'source_image': torch.rand([args.batch_size, 3, 240, 240], device=device), 'target_image': torch.rand([args.batch_size, 3, 240, 240], device=device), 'theta_GT': torch.rand([args.batch_size, 2, 3], device=device) } logs_writer.add_graph(model, dummy_input) # Start of training print('Starting training...') best_val_loss = float("inf") max_batch_iters = len(dataloader) print('Iterations for one epoch:', max_batch_iters) epoch_to_change_lr = int(args.lr_max_iter / max_batch_iters * 2 + 0.5) # Loading checkpoint model, optimizer, start_epoch, best_val_loss, last_epoch = load_checkpoint( checkpoint_path, model, optimizer, device) # Scheduler if args.lr_scheduler == 'cosine': is_cosine_scheduler = True scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.lr_max_iter, eta_min=1e-7, last_epoch=last_epoch) elif args.lr_scheduler == 'cosine_restarts': is_cosine_scheduler = True scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=args.lr_max_iter, T_mult=2, last_epoch=last_epoch) elif args.lr_scheduler == 'exp': is_cosine_scheduler = False if last_epoch > 0: last_epoch /= max_batch_iters scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=args.lr_decay, last_epoch=last_epoch) # elif args.lr_scheduler == 'step': # step_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 10, gamma=0.1) # scheduler = False else: is_cosine_scheduler = False scheduler = False for epoch in range(1, start_epoch): if args.lr_scheduler == 'cosine' and (epoch % epoch_to_change_lr == 0): scheduler.state_dict()['base_lrs'][0] *= args.lr_decay torch.autograd.set_detect_anomaly(True) for epoch in range(start_epoch, args.num_epochs + 1): print('Current epoch: ', epoch) # we don't need the average epoch loss so we assign it to _ _ = train(epoch, model, loss, optimizer, dataloader, pair_generation_tnf, log_interval=args.log_interval, scheduler=scheduler, is_cosine_scheduler=is_cosine_scheduler, tb_writer=logs_writer) # Step non-cosine scheduler if scheduler and not is_cosine_scheduler: scheduler.step() val_loss = validate_model(model, loss, dataloader_val, pair_generation_tnf, epoch, logs_writer) # Change lr_max in cosine annealing if args.lr_scheduler == 'cosine' and (epoch % epoch_to_change_lr == 0): scheduler.state_dict()['base_lrs'][0] *= args.lr_decay if (epoch % epoch_to_change_lr == epoch_to_change_lr // 2) or epoch == 1: compute_metric('absdiff', model, args.geometric_model, None, None, dataset_val, dataloader_val, pair_generation_tnf, args.batch_size, args) # remember best loss is_best = val_loss < best_val_loss best_val_loss = min(val_loss, best_val_loss) save_checkpoint( { 'epoch': epoch + 1, 'args': args, 'state_dict': model.state_dict(), 'best_val_loss': best_val_loss, 'optimizer': optimizer.state_dict(), }, is_best, checkpoint_path) logs_writer.close() print('Done!')
class CNNGeometricMatcher: def __init__(self, use_extracted_features=False, geometric_affine_model=None, geometric_tps_model=None, arch='resnet18', featext_weights=None, min_mutual_keypoints=6, min_reprojection_error=200): self.min_mutual_keypoints = min_mutual_keypoints self.min_reprojection_error = min_reprojection_error self.__do_affine = geometric_affine_model is not None self.__do_tps = not use_extracted_features and geometric_tps_model is not None self.__affTnf = GeometricTnf(geometric_model='affine', use_cuda=_use_cuda) if self.__do_affine: checkpoint = torch.load(geometric_affine_model, map_location=lambda storage, loc: storage) print('Loading CNN Affine Geometric Model') if use_extracted_features: self.model_affine = CNNGeometricRegression( use_cuda=use_cuda, geometric_model='affine', arch=arch, featext_weights=featext_weights) model_dict = self.model_affine.state_dict() pretrained_dict = { k: v for k, v in checkpoint['state_dict'].items() if k in model_dict } model_dict.update(pretrained_dict) self.model_affine.load_state_dict(model_dict) else: self.model_affine = CNNGeometric( use_cuda=_use_cuda, geometric_model='affine', arch=arch, featext_weights=featext_weights) self.model_affine.load_state_dict(checkpoint['state_dict']) self.model_affine.eval() if self.__do_tps: self.model_tps = CNNGeometric(use_cuda=use_cuda, geometric_model='tps', arch=arch, featext_weights=featext_weights) checkpoint = torch.load(geometric_tps_model, map_location=lambda storage, loc: storage) print('Loading CNN TPS Geometric Model') #self.model_tps.load_state_dict(checkpoint['state_dict']) model_dict = self.model_tps.state_dict() pretrained_dict = { k: v for k, v in checkpoint['state_dict'].items() if k in model_dict } model_dict.update(pretrained_dict) self.model_tps.load_state_dict(model_dict) self.model_tps.eval() self.pt = PointTnf(use_cuda=_use_cuda) def run(self, batch): if self.__do_affine: theta_aff, correlationAB, correlationBA = self.model_affine(batch) if self.__do_tps: if self.__do_affine: warped_image_aff = self.__affTnf(batch['source_image'], theta_aff.view(-1, 2, 3)) theta_affine_tps, correlationAB, correlationBA = self.model_tps( { 'source_image': warped_image_aff, 'target_image': batch['target_image'] }) else: theta_tps, correlationAB, correlationBA = self.model_tps(batch) keypoints_A, keypoints_B = find_mutual_matached_keypoints( correlationAB, correlationBA) num_mutual_keypoints = keypoints_A.shape[0] if num_mutual_keypoints < self.min_mutual_keypoints: matched = False reprojection_error = -1 else: source_im_size = batch['source_im_size'] target_im_size = batch['target_im_size'] source_im_shape_np = source_im_size.data.numpy() target_im_shape_np = target_im_size.data.numpy() tensor_shape = correlationAB.data.numpy()[0].shape im_keypointsA = tensorPointstoPixels( keypoints_A, tensor_size=tensor_shape, im_size=(source_im_shape_np[0][1], source_im_shape_np[0][0])) im_keypointsB = tensorPointstoPixels( keypoints_B, tensor_size=tensor_shape, im_size=(target_im_shape_np[0][1], target_im_shape_np[0][0])) torch_keypointsA_var = Variable( torch.Tensor( im_keypointsA.reshape(1, 2, -1).astype(np.float32))) torch_keypointsB_var = Variable( torch.Tensor( im_keypointsB.reshape(1, 2, -1).astype(np.float32))) target_points_norm = PointsToUnitCoords(torch_keypointsB_var, target_im_size) if self.__do_affine and self.__do_tps: warped_points_aff_tps_norm = self.pt.tpsPointTnf( theta_affine_tps, target_points_norm) warped_points_norm = self.pt.affPointTnf( theta_aff, warped_points_aff_tps_norm) elif self.__do_affine: warped_points_norm = self.pt.affPointTnf( theta_aff, target_points_norm) elif self.__do_tps: warped_points_norm = self.pt.tpsPointTnf( theta_tps, target_points_norm) warped_points_aff = PointsToPixelCoords(warped_points_norm, source_im_size) reprojection_error = compute_reprojection_error( torch_keypointsA_var, warped_points_aff) matched = reprojection_error <= self.min_reprojection_error return reprojection_error, matched, num_mutual_keypoints