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
