def load_model(aff_params_path = '', aff_feat_ext = 'wormbrain_1', aff_feat_reg = 'simpler', tps_params_path = '', tps_feat_ext = 'wormbrain_1', tps_feat_reg = 'simpler'): """ Loads a model. Assumes that each model (Affine and Thin-Plate Spline) have been trained separately. Must specify the architecture used for feature_extraction By default, it is resnet101 """ use_cuda = torch.cuda.is_available() #Only create a model for which weights have been provided do_aff = not aff_params_path=='' do_tps = not tps_params_path=='' if not do_aff and not do_tps: print("No weights found. Models not created, exiting.") return print("Creating CNN model.") if do_aff: model_aff = CNNGeometric(output_dim=6,use_cuda=use_cuda, feature_extraction_cnn= aff_feat_ext, feature_regression=aff_feat_reg) if do_tps: model_tps = CNNGeometric(output_dim=18,use_cuda=use_cuda, feature_extraction_cnn= tps_feat_ext, feature_regression = tps_feat_reg) print("Loading trained model weights.") if do_aff: #Loading affine model if aff_feat_ext == 'resnet101': aff_feat_ext = 'resnet' checkpoint = torch.load(aff_params_path, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([(k.replace(aff_feat_ext, 'model'), v) for k, v in checkpoint['state_dict'].items()]) model_aff.load_state_dict(checkpoint['state_dict']) print('Weights for Affine model loaded.') if do_tps: #Loading thin plate spline model if tps_feat_ext == 'resnet101': aff_feat_ext = 'resnet' checkpoint = torch.load(tps_params_path, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([(k.replace(aff_feat_ext, 'model'), v) for k, v in checkpoint['state_dict'].items()]) model_tps.load_state_dict(checkpoint['state_dict']) print('Weights for Thin-Plate Spline model loaded.') print('Returning model(s).') if do_aff and not do_tps: return model_aff if do_tps and not do_aff: return model_tps if do_aff and do_tps: return model_aff, model_tps
def create_model(model_filename): checkpoint = torch.load(model_filename, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()]) try: output_size = checkpoint['state_dict']['FeatureRegression.linear.bias'].size()[0] except: output_size = checkpoint['state_dict']['FeatureRegression.resnet.fc.bias'].size()[0] if output_size == 4: geometric_model = 'affine_simple_4' elif output_size == 3: geometric_model = 'affine_simple' else: raise NotImplementedError('Geometric model deducted from output layer is unsupported') model = CNNGeometric(use_cuda=use_cuda, output_dim=output_size, **arg_groups['model']) if use_me is False: for name, param in model.FeatureExtraction.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name]) for name, param in model.FeatureRegression.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureRegression.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression.' + name]) return (model,geometric_model)
def create_model(model_filename): checkpoint = torch.load(model_filename, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([ (k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items() ]) output_size = checkpoint['state_dict'][ 'FeatureRegression.linear.bias'].size()[0] if output_size == 6: geometric_model = 'affine' elif output_size == 8 or output_size == 9: geometric_model = 'hom' else: geometric_model = 'tps' model = CNNGeometric(use_cuda=use_cuda, output_dim=output_size, **arg_groups['model']) for name, param in model.FeatureExtraction.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureExtraction.state_dict()[name].copy_( checkpoint['state_dict']['FeatureExtraction.' + name]) for name, param in model.FeatureRegression.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureRegression.state_dict()[name].copy_( checkpoint['state_dict']['FeatureRegression.' + name]) return (model, geometric_model)
def load_torch_model(args, use_cuda): model = CNNGeometric(use_cuda=use_cuda, geometric_model='affine', feature_extraction_cnn=args.feature_extraction_cnn) # Load trained weights print('Loading trained model weights...') checkpoint = torch.load(args.pretrained, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict( [(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()]) model.load_state_dict(checkpoint['state_dict']) return model
# Download dataset if needed and set paths if args.training_dataset == 'pascal': if args.training_image_path == '': download_pascal('datasets/pascal-voc11/') args.training_image_path = 'datasets/pascal-voc11/' if args.training_tnf_csv == '' and args.geometric_model == 'affine': args.training_tnf_csv = 'training_data/pascal-synth-aff' elif args.training_tnf_csv == '' and args.geometric_model == 'tps': args.training_tnf_csv = 'training_data/pascal-synth-tps' # CNN model and loss print('Creating CNN model...') model = CNNGeometric(use_cuda=use_cuda, geometric_model=args.geometric_model, feature_extraction_cnn=args.feature_extraction_cnn) 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) # Dataset and dataloader dataset = SynthDataset(geometric_model=args.geometric_model, csv_file=os.path.join(args.training_tnf_csv, 'train.csv'), training_image_path=args.training_image_path,
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!')
tf.set_random_seed(args.seed) # Download dataset if needed and set paths if args.training_dataset == 'pascal': if args.training_image_path == '': download_pascal('datasets/pascal-voc11/') args.training_image_path = 'datasets/pascal-voc11/' if args.training_tnf_csv == '' and args.geometric_model == 'affine': args.training_tnf_csv = 'training_data/pascal-synth-aff' elif args.training_tnf_csv == '' and args.geometric_model == 'tps': args.training_tnf_csv = 'training_data/pascal-synth-tps' # CNN model and loss print('Creating CNN model...') model = CNNGeometric(feature_extraction_cnn=args.feature_extraction_cnn, training=False) if args.use_mse_loss: print('Using MSE loss...') loss = tf.losses else: print('Using grid loss...') loss = TransformedGridLoss(geometric_model=args.geometric_model) # Dataset and dataloader dataset = SynthDataset(geometric_model=args.geometric_model, csv_file=os.path.join(args.training_tnf_csv, 'train.csv'), training_image_path=args.training_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample)
use_cuda = torch.cuda.is_available() # Argument parsing parser = argparse.ArgumentParser(description='CNNGeometric PyTorch implementation') # Paths parser.add_argument('--model', type=str, default='trained_models/best_checkpoint_resnet18_adam_pose_mse_loss.pth.tar', help='Trained affine model filename') #parser.add_argument('--model-tps', type=str, default='trained_models/best_pascal_checkpoint_adam_tps_grid_loss.pth.tar', help='Trained TPS model filename') parser.add_argument('--path', type=str, default='/home/develop/Work/Datasets/', help='Path to PF dataset') parser.add_argument('--pairs', type=str, default='/home/develop/Work/Datasets/gardens_pairs_path_samples_sift_RANSAC_12kps.csv', help='Path to PF dataset') args = parser.parse_args() dataset_path=args.path dataset_pairs_file = args.pairs # Create model print('Creating CNN model...') model = CNNGeometric(use_cuda=use_cuda,geometric_model='pose',arch = 'resnet18') print('Load CNN Weights ...') checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage) model.load_state_dict(checkpoint['state_dict']) # Dataset and dataloader dataset = PlacesDataset(csv_file=dataset_pairs_file, training_image_path=dataset_path, transform=NormalizeImageDict(['source_image','target_image'])) dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4) batchTensorToVars = BatchTensorToVars(use_cuda=use_cuda) for source_im_path,target_im_path,batch in dataloader: # get random batch of size 1
print('Creating CNN model...') # check type of given model and create model # if two_stage: # model = TwoStageCNNGeometric(use_cuda=use_cuda, # **arg_groups['model']) # if not two_stage and do_aff: # model = CNNGeometric(use_cuda=use_cuda, # output_dim=6, # **arg_groups['model']) # if not two_stage and do_tps: # model_tps = CNNGeometric(use_cuda=use_cuda, # output_dim=18, # **arg_groups['model']) model = CNNGeometric(use_cuda=use_cuda, **arg_groups['model']) # load pretrained weights # if two_stage and args.model!='': # checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage) # checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()]) # for name, param in model.FeatureExtraction.state_dict().items(): # model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name]) # for name, param in model.FeatureRegression.state_dict().items(): # model.FeatureRegression.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression.' + name]) # for name, param in model.FeatureRegression2.state_dict().items(): # model.FeatureRegression2.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression2.' + name]) # for name, param in model.GridDeformation2.state_dict().items(): # model.GridDeformation2.state_dict()[name].copy_(checkpoint['state_dict']['GridDeformation2.' + name])
args = parser.parse_args() use_cuda = torch.cuda.is_available() do_aff = not args.model_aff == '' do_tps = not args.model_tps == '' # Download dataset if needed download_PF_willow('datasets/') # Create model print('Creating CNN model...') if do_aff: model_aff = CNNGeometric( use_cuda=use_cuda, geometric_model='affine', feature_extraction_cnn=args.feature_extraction_cnn) if do_tps: model_tps = CNNGeometric( use_cuda=use_cuda, geometric_model='tps', feature_extraction_cnn=args.feature_extraction_cnn) # Load trained weights print('Loading trained model weights...') if do_aff: checkpoint = torch.load(args.model_aff, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([ (k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()
torch.cuda.manual_seed(args.seed) # Download dataset if needed and set paths if args.training_dataset == 'pascal': if args.training_image_path == '': download_pascal('datasets/pascal-voc11/') args.training_image_path = 'datasets/pascal-voc11/' if args.training_tnf_csv == '' and args.geometric_model=='affine': args.training_tnf_csv = 'training_data/pascal-synth-aff' elif args.training_tnf_csv == '' and args.geometric_model=='tps': args.training_tnf_csv = 'training_data/pascal-synth-tps' # CNN model and loss print('Creating CNN model...') model = CNNGeometric(use_cuda=use_cuda,geometric_model=args.geometric_model) 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) # Dataset and dataloader dataset = SynthDataset(geometric_model=args.geometric_model, csv_file=os.path.join(args.training_tnf_csv,'train.csv'), training_image_path=args.training_image_path, transform=NormalizeImageDict(['image']))
def main(): args = parse_flags() use_cuda = torch.cuda.is_available() # 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.training_image_path == '': download_pascal('datasets/pascal-voc11/') args.training_image_path = 'datasets/pascal-voc11/' if args.training_tnf_csv == '' and args.geometric_model == 'affine': args.training_tnf_csv = 'training_data/pascal-synth-aff' elif args.training_tnf_csv == '' and args.geometric_model == 'tps': args.training_tnf_csv = 'training_data/pascal-synth-tps' # CNN model and loss if not args.pretrained: if args.light_model: print('Creating light CNN model...') model = LightCNN(use_cuda=use_cuda, geometric_model=args.geometric_model) else: print('Creating CNN model...') model = CNNGeometric( use_cuda=use_cuda, geometric_model=args.geometric_model, feature_extraction_cnn=args.feature_extraction_cnn) else: model = load_torch_model(args, use_cuda) if args.loss == 'mse': print('Using MSE loss...') loss = MSELoss() elif args.loss == 'sum': print('Using the sum of MSE and grid loss...') loss = GridLossWithMSE(use_cuda=use_cuda, geometric_model=args.geometric_model) else: print('Using grid loss...') loss = TransformedGridLoss(use_cuda=use_cuda, geometric_model=args.geometric_model) # Initialize csv paths train_csv_path_list = glob( os.path.join(args.training_tnf_csv, '*train.csv')) if len(train_csv_path_list) > 1: print( "!!!!WARNING!!!! multiple train csv files found, using first in glob order" ) elif not len(train_csv_path_list): raise FileNotFoundError( "No training csv where found in the specified path!!!") train_csv_path = train_csv_path_list[0] val_csv_path_list = glob(os.path.join(args.training_tnf_csv, '*val.csv')) if len(val_csv_path_list) > 1: print( "!!!!WARNING!!!! multiple train csv files found, using first in glob order" ) elif not len(val_csv_path_list): raise FileNotFoundError( "No training csv where found in the specified path!!!") val_csv_path = val_csv_path_list[0] # Initialize Dataset objects if args.coupled_dataset: # Dataset for train and val if dataset is already coupled dataset = CoupledDataset(geometric_model=args.geometric_model, csv_file=train_csv_path, training_image_path=args.training_image_path, transform=NormalizeImageDict( ['image_a', 'image_b'])) dataset_val = CoupledDataset( geometric_model=args.geometric_model, csv_file=val_csv_path, training_image_path=args.training_image_path, transform=NormalizeImageDict(['image_a', 'image_b'])) # Set Tnf pair generation func pair_generation_tnf = CoupledPairTnf(use_cuda=use_cuda) else: # Standard Dataset for train and val dataset = SynthDataset(geometric_model=args.geometric_model, csv_file=train_csv_path, training_image_path=args.training_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) dataset_val = SynthDataset( geometric_model=args.geometric_model, csv_file=val_csv_path, training_image_path=args.training_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 = Adam(model.FeatureRegression.parameters(), lr=args.lr) if args.lr_scheduler: if args.scheduler_type == 'cosine': print('Using cosine learning rate scheduler') scheduler = CosineAnnealingLR(optimizer, T_max=args.lr_max_iter, eta_min=args.lr_min) elif args.scheduler_type == 'decay': print('Using decay learning rate scheduler') scheduler = ReduceLROnPlateau(optimizer, 'min') else: print( 'Using truncated cosine with decay learning rate scheduler...') scheduler = TruncateCosineScheduler(optimizer, len(dataloader), args.num_epochs - 1) else: scheduler = False # Train # Set up names for checkpoints if args.loss == 'mse': ckpt = args.trained_models_fn + '_' + args.geometric_model + '_mse_loss' + args.feature_extraction_cnn checkpoint_path = os.path.join(args.trained_models_dir, args.trained_models_fn, ckpt + '.pth.tar') elif args.loss == 'sum': ckpt = args.trained_models_fn + '_' + args.geometric_model + '_sum_loss' + args.feature_extraction_cnn checkpoint_path = os.path.join(args.trained_models_dir, args.trained_models_fn, ckpt + '.pth.tar') else: ckpt = args.trained_models_fn + '_' + args.geometric_model + '_grid_loss' + args.feature_extraction_cnn checkpoint_path = os.path.join(args.trained_models_dir, args.trained_models_fn, ckpt + '.pth.tar') if not os.path.exists(args.trained_models_dir): os.mkdir(args.trained_models_dir) # Set up TensorBoard writer if not args.log_dir: tb_dir = os.path.join(args.trained_models_dir, args.trained_models_fn + '_tb_logs') else: tb_dir = os.path.join(args.log_dir, args.trained_models_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]), 'target_image': torch.rand([args.batch_size, 3, 240, 240]), 'theta_GT': torch.rand([16, 2, 3]) } 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, coupled=args.coupled_dataset) # 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!')
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)
training_image_path=args.pf_path, transform=NormalizeImageDict( ['source_image', 'target_image'])) batchTensorToVars = BatchTensorToVars() # Instantiate point transformer pt = PointTnf() # Instatiate image transformers affTnf = GeometricTnf(geometric_model='affine') with tf.Graph().as_default(): # Create model print('Creating CNN model...') model_aff = CNNGeometric( feature_extraction_cnn=args.feature_extraction_cnn) init_op = tf.global_variables_initializer() saver = tf.train.Saver() with tf.Session() as sess: sess.run(init_op) # Load trained weights print('Loading trained model weights...') saver.restore(sess, args.model_aff) # reload_checkpoint(sess, saver, args.model_aff) while True: # get random batch of size 1 batch = choice(dataset) batch = batchTensorToVars(batch)
args.dataset_csv_path = 'training_data/pascal-synth-aff' elif args.dataset_csv_path == '' and args.geometric_model == 'tps': args.dataset_csv_path = 'training_data/pascal-synth-tps' arg_groups['dataset']['dataset_image_path'] = args.dataset_image_path # CNN model and loss print('Creating CNN model...') if args.geometric_model == 'affine': cnn_output_dim = 6 elif args.geometric_model == 'tps': cnn_output_dim = 256 # cnn_output_dim = 256 model = CNNGeometric(use_cuda=use_cuda, output_dim=cnn_output_dim, **arg_groups['model']) # Load pretrained model if args.model != '': checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([ (k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items() ]) for name, param in model.FeatureExtraction.state_dict().items(): model.FeatureExtraction.state_dict()[name].copy_( checkpoint['state_dict']['FeatureExtraction.' + name]) for name, param in model.FeatureRegression.state_dict().items(): model.FeatureRegression.state_dict()[name].copy_(
args = parser.parse_args() use_cuda = torch.cuda.is_available() do_aff = not args.model_aff == '' do_tps = not args.model_tps == '' # Download dataset if needed download_PF_willow('datasets/') # Create model print('Creating CNN model...') if do_aff: model_aff = CNNGeometric( use_cuda=use_cuda, output_dim=6, feature_extraction_cnn=args.feature_extraction_cnn) if do_tps: model_tps = CNNGeometric( use_cuda=use_cuda, output_dim=18, feature_extraction_cnn=args.feature_extraction_cnn) # Load trained weights print('Loading trained model weights...') if do_aff: checkpoint = torch.load(args.model_aff, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([ (k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()
elif args.eval_dataset=='caltech' and not exists(args.eval_dataset_path): download_caltech(args.eval_dataset_path) elif args.eval_dataset=='tss' and not exists(args.eval_dataset_path): download_TSS(args.eval_dataset_path) elif args.eval_dataset=='pascal-parts' and not exists(args.eval_dataset_path): download_pascal_parts(args.eval_dataset_path) # Create model print('Creating CNN model...') # check type of given model and create model if two_stage: model = TwoStageCNNGeometric(use_cuda=use_cuda, **arg_groups['model']) if not two_stage and do_aff: model = CNNGeometric(use_cuda=use_cuda, output_dim=6, **arg_groups['model']) if not two_stage and do_tps: model_tps = CNNGeometric(use_cuda=use_cuda, output_dim=18, **arg_groups['model']) # load pretrained weights if two_stage and args.model!='': checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()]) for name, param in model.FeatureExtraction.state_dict().items(): model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name]) for name, param in model.FeatureRegression.state_dict().items():
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!')
default='Pairs CSV file', help='Path to PF dataset') args = parser.parse_args() use_cuda = torch.cuda.is_available() do_aff = not args.model_aff == '' do_tps = False #not args.model_tps=='' show_diffmap = True dataset_path = args.path dataset_pairs_file = args.pairs # Create model print('Creating CNN model...') if do_aff: model_aff = CNNGeometric(use_cuda=use_cuda, geometric_model='affine') if do_tps: model_tps = CNNGeometric(use_cuda=use_cuda, geometric_model='tps') # Load trained weights print('Loading trained model weights...') if do_aff: checkpoint = torch.load(args.model_aff, map_location=lambda storage, loc: storage) model_aff.load_state_dict(checkpoint['state_dict']) if do_tps: checkpoint = torch.load(args.model_tps, map_location=lambda storage, loc: storage) model_tps.load_state_dict(checkpoint['state_dict']) # Dataset and dataloader
torch.cuda.manual_seed(args.seed) # Download dataset if needed and set paths if args.training_dataset == 'pascal': if args.training_image_path == '': download_pascal('datasets/pascal-voc11/') args.training_image_path = 'datasets/pascal-voc11/' if args.training_tnf_csv == '' and args.geometric_model=='affine': args.training_tnf_csv = 'training_data/pascal-synth-aff' elif args.training_tnf_csv == '' and args.geometric_model=='tps': args.training_tnf_csv = 'training_data/pascal-synth-tps' # CNN model and loss print('Creating CNN model...') model = CNNGeometric(use_cuda=use_cuda,geometric_model=args.geometric_model,arch=args.arch,featext_weights=args.fw) 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) if args.geometric_model == 'pose': dataset = PoseDataset( geometric_model=args.geometric_model, csv_file=os.path.join(args.training_tnf_csv, 'train.csv'), training_image_path=args.training_image_path, output_size=(240, 240),