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, transform=NormalizeImageDict(['image']), dataset_csv_file='train.csv', **arg_groups['dataset']) dataloader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=4) # don't change num_workers, as they copy the rnd seed dataset_test = SynthDataset(geometric_model=args.geometric_model, transform=NormalizeImageDict(['image']), dataset_csv_file='test.csv', **arg_groups['dataset']) dataloader_test = DataLoader(dataset_test, batch_size=args.batch_size, shuffle=True, num_workers=4) cnn_image_size = (args.image_size, args.image_size) pair_generation_tnf = SynthPairTnf(geometric_model=args.geometric_model,
# Download validation dataset if needed if args.eval_dataset_path == '' and args.eval_dataset == 'pf-pascal': args.eval_dataset_path = 'datasets/proposal-flow-pascal/' if args.eval_dataset == 'pf-pascal' and not exists(args.eval_dataset_path): download_PF_pascal(args.eval_dataset_path) batch_tnf = BatchTensorToVars(use_cuda=use_cuda) dataset_test = PFPascalDataset( csv_file=os.path.join(args.eval_dataset_path, 'test_pairs_pf_pascal.csv'), dataset_path=args.eval_dataset_path, transform=NormalizeImageDict(['source_image', 'target_image'])) dataloader_test = DataLoader(dataset_test, batch_size=compute_metric_batch, shuffle=False, num_workers=4) model.load_state_dict(torch.load("./trained_models/jim_net.pth.tar")) model.eval() print('Final test pck...') compute_metric('pck', model.Alignment_Module, dataset_test, dataloader_test, batch_tnf, compute_metric_batch, True, True, True, args) true_labels = [] predicted_labels = [] saved_image_clusters = {}
# filter training categories if args.categories != 0: keep = np.zeros((len(dataset.set), 1)) for i in range(len(dataset.set)): keep[i] = np.sum(dataset.set[i] == args.categories) keep_idx = np.nonzero(keep)[0] dataset.set = dataset.set[keep_idx] dataset.img_A_names = dataset.img_A_names[keep_idx] dataset.img_B_names = dataset.img_B_names[keep_idx] batch_tnf = BatchTensorToVars(use_cuda=use_cuda) # dataloader dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4) dataloader_eval = DataLoader(dataset_eval, batch_size=args.batch_size, shuffle=False, num_workers=4) checkpoint_name = os.path.join( args.result_model_dir, 'inlier_{}_consis_{}_trans_{}.pth.tar'.format(args.w_inlier, args.w_consis, args.w_trans)) # define epoch function def process_epoch(mode, epoch,
def process_epoch(mode, epoch, model, loss_fn, optimizer, dataloader, batch_preprocessing_fn, use_cuda=True, log_interval=50): for param_group in warm_optimizer.param_groups: if epoch > 1: param_group['lr'] = args.classifier_lr else: param_group['lr'] = 0. for i, param_group in enumerate(optimizer.param_groups): if epoch > 1: if i == 0: param_group['lr'] = args.align_lr elif i <= 3: param_group['lr'] = args.classifier_lr else: raise NotImplementedError else: param_group['lr'] = 0. print_time("Started Feature Extraction Step: " + str(epoch)) old_clusters = cluster_class.clusters #cluster step cluster_features = [] images = [] true_labels = [] with torch.no_grad(): for batch_idx, batch in enumerate(dataloader): tnf_batch = batch_preprocessing_fn(batch) _, _, _, _, _, cluster_feature = model(tnf_batch) cluster_features.append(cluster_feature.cpu().detach().numpy()) images.append(tnf_batch['source_image'].cpu().detach().numpy()) #TODO make this an optional flag #only used for computing metrics true_labels.append(tnf_batch['set'].cpu().detach().numpy()) with torch.no_grad(): for batch_idx, batch in enumerate(dataloader): tnf_batch = batch_preprocessing_fn(batch) _, _, _, _, _, cluster_feature = model(tnf_batch) cluster_features.append(cluster_feature.cpu().detach().numpy()) images.append(tnf_batch['target_image'].cpu().detach().numpy()) #TODO make this an optional flag #only used for computing metrics true_labels.append(tnf_batch['set'].cpu().detach().numpy()) cluster_features = np.concatenate(cluster_features) images = np.concatenate(images) true_labels = np.concatenate(true_labels) print_time("Started Clustering: " + str(epoch)) tsne = TSNE(n_components=3, n_jobs=40, init='pca', perplexity=25) features_tsne = tsne.fit_transform(cluster_features) gmm = GaussianMixture(n_components=num_clusters, covariance_type='full') gmm.fit(features_tsne) clusters = gmm.predict(features_tsne) if len(old_clusters) != 0: aligned_clusters = align_cluster_index(old_clusters, clusters) clusters = aligned_clusters cluster_class.clusters = clusters probabilities = one_hot(clusters, num_clusters) print_time("Started Warm Start: " + str(epoch)) warm_start_tensor = TensorDataset(torch.Tensor(cluster_features), torch.Tensor(probabilities)) warm_start_loader = DataLoader(warm_start_tensor, batch_size=args.batch_size, shuffle=True) loss_warm_start = 0. for i in range(args.warmup_epochs): for batch_idx, (batch, label) in enumerate(warm_start_loader): if mode == 'train': optimizer.zero_grad() _, loss = model.classify(batch.cuda(), probabilities=label.cuda()) loss_warm_start += loss.data.cpu().numpy() if mode == 'train': loss.backward() warm_optimizer.step() print_time("Warm start loss: " + str(loss_warm_start / (args.warmup_epochs * len(warm_start_loader) + eps))) print_time("Started Alignment: " + str(epoch)) #alignment step unsupervised_dataset = UnsupervisedDataset(images, cluster_features, clusters, probabilities, num_clusters) unsupervised_dataloader = DataLoader(unsupervised_dataset, batch_size=args.batch_size, shuffle=True) epoch_loss = 0 for batch_idx, batch in enumerate(unsupervised_dataloader): if mode == 'train': optimizer.zero_grad() tnf_batch = batch_preprocessing_fn(batch) loss = loss_fn(tnf_batch) loss_np = loss.data.cpu().numpy() epoch_loss += loss_np if mode == 'train': loss.backward() optimizer.step() else: loss = None if batch_idx % log_interval == 0: print( mode.capitalize() + ' Epoch: {} [{}/{} ({:.0f}%)]\t\tLoss: {:.6f}'.format( epoch, batch_idx, len(unsupervised_dataloader), 100. * batch_idx / len(unsupervised_dataloader), loss_np)) epoch_loss /= len(unsupervised_dataloader) print_time(mode.capitalize() + ' set: Average loss: {:.4f}'.format(epoch_loss)) #evaluate classifier cluster_accuracy_loader = DataLoader(torch.Tensor(images), batch_size=args.batch_size, shuffle=False) predicted_cluster = [] with torch.no_grad(): for batch_idx, batch in enumerate(cluster_accuracy_loader): preds = model.classify(batch.cuda(), processed=False) predicted_cluster.append(preds.cpu().detach().numpy()) pred_clusters = np.concatenate(predicted_cluster).reshape( images.shape[0], num_clusters) pred_clusters = np.argmax(pred_clusters, axis=1) print_time( "V measure predictions: " + str(sklearn.metrics.v_measure_score(true_labels, pred_clusters))) print_time("V measure clusters: " + str(sklearn.metrics.v_measure_score(true_labels, clusters))) try: aligned = align_cluster_index(true_labels - 1, pred_clusters) equality = np.equal(aligned, true_labels - 1) aligned_with_true = equality.sum() / len(true_labels) clust_acc1.append((epoch, aligned_with_true)) print_time("Accuracy Predictions: " + str(aligned_with_true)) except: print_time("Failed to align predictions with labels") try: aligned = align_cluster_index(true_labels - 1, clusters) equality = np.equal(aligned, true_labels - 1) aligned_with_true = equality.sum() / len(true_labels) clust_acc2.append((epoch, aligned_with_true)) print_time("Accuracy Clustering: " + str(aligned_with_true)) except: print_time("Failed to align clusters with labels") return epoch_loss