def main(): cv2.setNumThreads(1) # Retrieve config file p = create_config(args.config_env, args.config_exp) print('Python script is {}'.format(os.path.abspath(__file__))) print(colored(p, 'red')) # Get model print(colored('Retrieve model', 'blue')) model = get_model(p) print(model) model = model.cuda() # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'blue')) # Transforms val_transforms = get_val_transformations() val_dataset = get_val_dataset(p, val_transforms) true_val_dataset = get_val_dataset(p, None) # True validation dataset without reshape val_dataloader = get_val_dataloader(p, val_dataset) print(colored('Val samples %d' %(len(val_dataset)), 'yellow')) # Evaluate best model at the end print(colored('Evaluating model at {}'.format(args.state_dict), 'blue')) model.load_state_dict(torch.load(args.state_dict, map_location='cpu')) save_results_to_disk(p, val_dataloader, model, crf_postprocess=args.crf_postprocess) eval_stats = eval_segmentation_supervised_offline(p, true_val_dataset, verbose=True)
def main(): cv2.setNumThreads(1) # Retrieve config file p = create_config(args.config_env, args.config_exp) print('Python script is {}'.format(os.path.abspath(__file__))) print(colored(p, 'red')) # Get model print(colored('Retrieve model', 'blue')) model = get_model(p) print(model) model = model.cuda() # Load pre-trained weights state_dict = torch.load(p['pretraining'], map_location='cpu') # State dict follows our lay-out if 'model' in state_dict.keys(): state_dict = state_dict['model'] new_state = {} for k, v in state_dict.items(): if k.startswith('module.model_q'): new_state[k.rsplit('module.model_q.')[1]] = v msg = model.load_state_dict(new_state, strict=False) print(msg) # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'blue')) # Transforms from data.dataloaders.pascal_voc import VOC12 val_transforms = get_val_transformations() print(val_transforms) val_dataset = VOC12(split='val', transform=val_transforms) val_dataloader = get_val_dataloader(p, val_dataset) true_val_dataset = VOC12(split='val', transform=None) print(colored('Val samples %d' %(len(true_val_dataset)), 'yellow')) # Kmeans Clustering n_clusters = 21 results_miou = [] for i in range(args.num_seeds): save_embeddings_to_disk(p, val_dataloader, model, n_clusters=n_clusters, seed=1234 + i) eval_stats = eval_kmeans(p, true_val_dataset, n_clusters=n_clusters, verbose=True) results_miou.append(eval_stats['mIoU']) print(colored('Average mIoU is %2.1f' %(np.mean(results_miou)*100), 'green'))
def main(): args = FLAGS.parse_args() p = create_config(args.config_env, args.config_exp, args.tb_run) print(colored(p, 'red')) # CUDNN torch.backends.cudnn.benchmark = True # Data print(colored('Get dataset and dataloaders', 'blue')) train_transformations = get_train_transformations(p) val_transformations = get_val_transformations(p) train_dataset = get_train_dataset(p, train_transformations, split='train', to_similarity_dataset=True) val_dataset = get_val_dataset(p, val_transformations, to_similarity_dataset=True) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Train transforms:', train_transformations) print('Validation transforms:', val_transformations) print('Train samples %d - Val samples %d' % (len(train_dataset), len(val_dataset))) # Tensorboard writer writer = SummaryWriter(log_dir=p['simpred_tb_dir']) # Model print(colored('Get model', 'blue')) model = get_model(p, p['pretext_model']) print(model) model = torch.nn.DataParallel(model) model = model.cuda() # Optimizer print(colored('Get optimizer', 'blue')) optimizer = get_optimizer(p, model, p['update_cluster_head_only']) print(optimizer) # Warning if p['update_cluster_head_only']: print(colored('WARNING: will only update the cluster head', 'red')) # Loss function print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() print(criterion) # Checkpoint if os.path.exists(p['simpred_checkpoint']): print(colored('Restart from checkpoint {}'.format(p['simpred_checkpoint']), 'blue')) checkpoint = torch.load(p['simpred_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] best_acc = checkpoint['best_acc'] else: print(colored('No checkpoint file at {}'.format(p['simpred_checkpoint']), 'blue')) start_epoch = 0 best_acc = 0 # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' % (epoch + 1, p['epochs']), 'yellow')) print(colored('-' * 15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') simpred_train(train_dataloader, model, criterion, optimizer, epoch, writer, p['update_cluster_head_only']) # Evaluate print('Make prediction on validation set ...') predictions = get_predictions(p, val_dataloader, model) print('Evaluate based on simpred loss ...') simpred_stats = simpred_evaluate(predictions, writer, epoch) print(simpred_stats) accuracy = simpred_stats['accuracy'] if accuracy > best_acc: print('New highest accuracy on validation set: %.4f -> %.4f' % (best_acc, accuracy)) best_acc = accuracy torch.save({'model': model.module.state_dict()}, p['simpred_model']) else: print('No new highest accuracy on validation set: %.4f -> %.4f' % (best_acc, accuracy)) # Checkpoint print('Checkpoint ...') torch.save({'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1, 'best_acc': best_acc}, p['simpred_checkpoint']) # Evaluate and save the final model print(colored('Evaluate best model based on simpred metric at the end', 'blue')) model_checkpoint = torch.load(p['simpred_model'], map_location='cpu') model.module.load_state_dict(model_checkpoint['model']) predictions, features, thumbnails = get_predictions(p, val_dataloader, model, return_features=True, return_thumbnails=True) writer.add_embedding(features, predictions[0]['targets'], thumbnails, p['epochs'], p['simpred_tb_dir'])
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # Get model print(colored('Retrieve model', 'blue')) model = get_model(p, p['scan_model']) print(model) model = torch.nn.DataParallel(model) model = model.cuda() # Get criterion print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() print(criterion) # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Optimizer print(colored('Retrieve optimizer', 'blue')) optimizer = get_optimizer(p, model) print(optimizer) # Dataset print(colored('Retrieve dataset', 'blue')) # Transforms strong_transforms = get_train_transformations(p) val_transforms = get_val_transformations(p) train_dataset = get_train_dataset(p, {'standard': val_transforms, 'augment': strong_transforms}, split='train', to_augmented_dataset=True) train_dataloader = get_train_dataloader(p, train_dataset) val_dataset = get_val_dataset(p, val_transforms) val_dataloader = get_val_dataloader(p, val_dataset) print(colored('Train samples %d - Val samples %d' %(len(train_dataset), len(val_dataset)), 'yellow')) # Checkpoint if os.path.exists(p['selflabel_checkpoint']): print(colored('Restart from checkpoint {}'.format(p['selflabel_checkpoint']), 'blue')) checkpoint = torch.load(p['selflabel_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] else: print(colored('No checkpoint file at {}'.format(p['selflabel_checkpoint']), 'blue')) start_epoch = 0 # EMA if p['use_ema']: ema = EMA(model, alpha=p['ema_alpha']) else: ema = None # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' %(epoch+1, p['epochs']), 'yellow')) print(colored('-'*10, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Perform self-labeling print('Train ...') selflabel_train(train_dataloader, model, criterion, optimizer, epoch, ema=ema) # Evaluate (To monitor progress - Not for validation) print('Evaluate ...') predictions = get_predictions(p, val_dataloader, model) clustering_stats = hungarian_evaluate(0, predictions, compute_confusion_matrix=False) print(clustering_stats) # Checkpoint print('Checkpoint ...') torch.save({'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1}, p['selflabel_checkpoint']) #torch.save(model.module.state_dict(), p['selflabel_model']) torch.save(model.module.state_dict(), os.path.join(p['selflabel_dir'], 'model_%d.pth.tar' %(epoch))) # Evaluate and save the final model print(colored('Evaluate model at the end', 'blue')) predictions = get_predictions(p, val_dataloader, model) clustering_stats = hungarian_evaluate(0, predictions, class_names=val_dataset.classes, compute_confusion_matrix=True, confusion_matrix_file=os.path.join(p['selflabel_dir'], 'confusion_matrix.png')) print(clustering_stats) torch.save(model.module.state_dict(), p['selflabel_model'])
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # Model print(colored('Retrieve model', 'blue')) model = get_model(p) print('Model is {}'.format(model.__class__.__name__)) print('Model parameters: {:.2f}M'.format(sum(p.numel() for p in model.parameters()) / 1e6)) print(model) model = model.cuda() # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'blue')) train_transforms = get_train_transformations(p) print('Train transforms:', train_transforms) val_transforms = get_val_transformations(p) print('Validation transforms:', val_transforms) train_dataset = get_train_dataset(p, train_transforms, to_augmented_dataset=True, split='train+unlabeled') # Split is for stl-10 val_dataset = get_val_dataset(p, val_transforms) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Dataset contains {}/{} train/val samples'.format(len(train_dataset), len(val_dataset))) # Memory Bank print(colored('Build MemoryBank', 'blue')) base_dataset = get_train_dataset(p, val_transforms, split='train') # Dataset w/o augs for knn eval base_dataloader = get_val_dataloader(p, base_dataset) memory_bank_base = MemoryBank(len(base_dataset), p['model_kwargs']['features_dim'], p['num_classes'], p['criterion_kwargs']['temperature']) memory_bank_base.cuda() memory_bank_val = MemoryBank(len(val_dataset), p['model_kwargs']['features_dim'], p['num_classes'], p['criterion_kwargs']['temperature']) memory_bank_val.cuda() # Criterion print(colored('Retrieve criterion', 'blue')) criterion = get_criterion(p) print('Criterion is {}'.format(criterion.__class__.__name__)) criterion = criterion.cuda() # Optimizer and scheduler print(colored('Retrieve optimizer', 'blue')) optimizer = get_optimizer(p, model) print(optimizer) # Checkpoint if os.path.exists(p['pretext_checkpoint']): print(colored('Restart from checkpoint {}'.format(p['pretext_checkpoint']), 'blue')) checkpoint = torch.load(p['pretext_checkpoint'], map_location='cpu') optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['model']) model.cuda() start_epoch = checkpoint['epoch'] else: print(colored('No checkpoint file at {}'.format(p['pretext_checkpoint']), 'blue')) start_epoch = 0 model = model.cuda() # Training print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' %(epoch, p['epochs']), 'yellow')) print(colored('-'*15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') simclr_train(train_dataloader, model, criterion, optimizer, epoch) # Fill memory bank print('Fill memory bank for kNN...') fill_memory_bank(base_dataloader, model, memory_bank_base) # Evaluate (To monitor progress - Not for validation) print('Evaluate ...') top1 = contrastive_evaluate(val_dataloader, model, memory_bank_base) print('Result of kNN evaluation is %.2f' %(top1)) # Checkpoint print('Checkpoint ...') torch.save({'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1}, p['pretext_checkpoint']) if epoch in [50, 75]: # Save final model # torch.save(model.state_dict(), p['pretext_model']) # Mine the topk nearest neighbors at the very end (Train) # These will be served as input to the SCAN loss. print(colored('Fill memory bank for mining the nearest neighbors (train) ...', 'blue')) fill_memory_bank(base_dataloader, model, memory_bank_base) topk = 20 print('Mine the nearest neighbors (Top-%d)' %(topk)) indices, acc = memory_bank_base.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' %(topk, 100*acc)) # np.save(p['topk_neighbors_train_path'], indices) # Mine the topk nearest neighbors at the very end (Val) # These will be used for validation. print(colored('Fill memory bank for mining the nearest neighbors (val) ...', 'blue')) fill_memory_bank(val_dataloader, model, memory_bank_val) topk = 5 print('Mine the nearest neighbors (Top-%d)' %(topk)) indices, acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' %(topk, 100*acc)) # np.save(p['topk_neighbors_val_path'], indices) # Save final model torch.save(model.state_dict(), p['pretext_model']) # Mine the topk nearest neighbors at the very end (Train) # These will be served as input to the SCAN loss. print(colored('Fill memory bank for mining the nearest neighbors (train) ...', 'blue')) fill_memory_bank(base_dataloader, model, memory_bank_base) topk = 20 print('Mine the nearest neighbors (Top-%d)' %(topk)) indices, acc = memory_bank_base.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' %(topk, 100*acc)) np.save(p['topk_neighbors_train_path'], indices) # Mine the topk nearest neighbors at the very end (Val) # These will be used for validation. print(colored('Fill memory bank for mining the nearest neighbors (val) ...', 'blue')) fill_memory_bank(val_dataloader, model, memory_bank_val) topk = 5 print('Mine the nearest neighbors (Top-%d)' %(topk)) indices, acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' %(topk, 100*acc)) np.save(p['topk_neighbors_val_path'], indices)
def main(): #try: # Retrieve config file cv2.setNumThreads(0) p = create_config(args.config_env, args.config_exp, args.save_name) sys.stdout = Logger(os.path.join(p['output_dir'], 'log_file.txt')) print(colored(p, 'red')) # Get model print(colored('Retrieve model', 'blue')) model = get_model(p) model = torch.nn.DataParallel(model) model = model.cuda() # device=device) # Get criterion print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() # device=device) print(criterion) # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Optimizer print(colored('Retrieve optimizer', 'blue')) optimizer = get_optimizer(p, model) print(optimizer) # Dataset print(colored('Retrieve dataset', 'blue')) # Transforms train_transforms, val_transforms = get_transformations(p) train_dataset = get_train_dataset(p, train_transforms) val_dataset = get_val_dataset(p, val_transforms) true_val_dataset = get_val_dataset( p, None) # True validation dataset without reshape train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Train samples %d - Val samples %d' % (len(train_dataset), len(val_dataset))) print('Train transformations:') print(train_transforms) print('Val transformations:') print(val_transforms) # Resume from checkpoint if os.path.exists(p['checkpoint']): print( colored('Restart from checkpoint {}'.format(p['checkpoint']), 'blue')) checkpoint = torch.load(p['checkpoint'], map_location='cpu') optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['model']) start_epoch = checkpoint['epoch'] best_result = checkpoint['best_result'] else: print( colored('No checkpoint file at {}'.format(p['checkpoint']), 'blue')) start_epoch = 0 save_model_predictions(p, val_dataloader, model) best_result = eval_all_results(p) # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' % (epoch + 1, p['epochs']), 'yellow')) print(colored('-' * 10, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') eval_train = train_vanilla(p, train_dataloader, model, criterion, optimizer, epoch) # Evaluate # Check if need to perform eval first if 'eval_final_10_epochs_only' in p.keys( ) and p['eval_final_10_epochs_only']: # To speed up -> Avoid eval every epoch, and only test during final 10 epochs. if epoch + 1 > p['epochs'] - 10: eval_bool = True else: eval_bool = False else: eval_bool = True # Perform evaluation if eval_bool: print('Evaluate ...') save_model_predictions(p, val_dataloader, model) curr_result = eval_all_results(p) improves, best_result = validate_results(p, curr_result, best_result) if improves: print('Save new best model') torch.save(model.state_dict(), p['best_model']) # Checkpoint print('Checkpoint ...') torch.save( { 'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1, 'best_result': best_result }, p['checkpoint']) # Evaluate best model at the end print(colored('Evaluating best model at the end', 'blue')) model.load_state_dict(torch.load(p['checkpoint'])['model']) print("Model state dict keys: ", model.state_dict().keys()) #print("Model state dict all: ", model.state_dict().items()) save_model_predictions(p, val_dataloader, model) eval_stats = eval_all_results(p) send_email(target_mail_address_list, server_name=server_name, exception_message="Success!", successfully=True)
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # Model print(colored('Retrieve model', 'green')) model = get_model(p) print('Model is {}'.format(model.__class__.__name__)) print('Model parameters: {:.2f}M'.format( sum(p.numel() for p in model.parameters()) / 1e6)) print(model) model = model.to(device) # CUDNN print(colored('Set CuDNN benchmark', 'green')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'green')) train_transforms = get_train_transformations(p) print('Train transforms:', train_transforms) val_transforms = get_val_transformations(p) print('Validation transforms:', val_transforms) train_dataset = get_train_dataset(p, train_transforms, to_augmented_dataset=True, split='train') # Split is for stl-10 val_dataset = get_val_dataset(p, val_transforms) train_dataloader = get_val_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Dataset contains {}/{} train/val samples'.format( len(train_dataset), len(val_dataset))) # Memory Bank print(colored('Build MemoryBank', 'green')) base_dataset = get_train_dataset( p, val_transforms, split='train') # Dataset w/o augs for knn eval base_dataloader = get_val_dataloader(p, base_dataset) memory_bank_base = MemoryBank(len(base_dataset), p['model_kwargs']['features_dim'], p['num_classes'], p['criterion_kwargs']['temperature']) memory_bank_base.to(device) memory_bank_val = MemoryBank(len(val_dataset), p['model_kwargs']['features_dim'], p['num_classes'], p['criterion_kwargs']['temperature']) memory_bank_val.to(device) # Checkpoint if os.path.exists(p['pretext_checkpoint']): print( colored( 'Restart from checkpoint {}'.format(p['pretext_checkpoint']), 'green')) checkpoint = torch.load(p['pretext_checkpoint'], map_location='cpu') # optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['model']) model.to(device) # start_epoch = checkpoint['epoch'] else: print( colored('No checkpoint file at {}'.format(p['pretext_checkpoint']), 'green')) start_epoch = 0 model = model.to(device) # # Training # print(colored('Starting main loop', 'green')) # with torch.no_grad(): # model.eval() # total_top1, total_top5, total_num, feature_bank = 0.0, 0.0, 0, [] # # # progress_bar = tqdm(train_dataloader) # for idx, batch in enumerate(train_dataloader): # images = batch['image'].to(device, non_blocking=True) # # target = batch['target'].to(device, non_blocking=True) # # output = model(images) # feature = F.normalize(output, dim=1) # feature_bank.append(feature) # # if idx % 25 == 0: # print("Feature bank buidling : {} / {}".format(idx, len(train_dataset)/p["batch_size"])) # # # [D, N] # feature_bank = torch.cat(feature_bank, dim=0).t().contiguous() # print(colored("Feature bank created. Similarity index starts now", "green")) # print(feature_bank.size()) # # for idx, batch in enumerate(train_dataloader): # # images = batch['image'].to(device, non_blocking=True) # # target = batch['target'].to(device, non_blocking=True) # # output = model(images) # feature = F.normalize(output, dim=1) # # sim_indices = knn_predict(feature, feature_bank, "", "", 10, 0.1) # # print(sim_indices) # # if idx == 10: # break # # Mine the topk nearest neighbors at the very end (Train) # # These will be served as input to the SCAN loss. # print(colored('Fill memory bank for mining the nearest neighbors (train) ...', 'green')) # fill_memory_bank(base_dataloader, model, memory_bank_base) # topk = 20 # print('Mine the nearest neighbors (Top-%d)' %(topk)) # indices, acc = memory_bank_base.mine_nearest_neighbors(topk) # print('Accuracy of top-%d nearest neighbors on train set is %.2f' %(topk, 100*acc)) # np.save(p['topk_neighbors_train_path'], indices) # Mine the topk nearest neighbors at the very end (Val) # These will be used for validation. print( colored('Fill memory bank for mining the nearest neighbors (val) ...', 'green')) fill_memory_bank(val_dataloader, model, memory_bank_val) topk = 5 print('Mine the nearest neighbors (Top-%d)' % (topk)) indices, acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' % (topk, 100 * acc)) np.save(p['topk_neighbors_val_path'], indices)
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # Model print(colored('Retrieve model', 'blue')) model = get_model(p) print('Model is {}'.format(model.__class__.__name__)) print('Model parameters: {:.2f}M'.format( sum(p.numel() for p in model.parameters()) / 1e6)) print(model) model = model.cuda() # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Dataset val_transforms = get_val_transformations(p) print('Validation transforms:', val_transforms) val_dataset = get_val_dataset(p, val_transforms) val_dataloader = get_val_dataloader(p, val_dataset) print('Dataset contains {} val samples'.format(len(val_dataset))) # Memory Bank print(colored('Build MemoryBank', 'blue')) base_dataset = get_train_dataset( p, val_transforms, split='train') # Dataset w/o augs for knn eval base_dataloader = get_val_dataloader(p, base_dataset) memory_bank_base = MemoryBank(len(base_dataset), p['model_kwargs']['features_dim'], p['num_classes'], p['criterion_kwargs']['temperature']) memory_bank_base.cuda() memory_bank_val = MemoryBank(len(val_dataset), p['model_kwargs']['features_dim'], p['num_classes'], p['criterion_kwargs']['temperature']) memory_bank_val.cuda() # Checkpoint assert os.path.exists(p['pretext_checkpoint']) print( colored('Restart from checkpoint {}'.format(p['pretext_checkpoint']), 'blue')) checkpoint = torch.load(p['pretext_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint) model.cuda() # Save model torch.save(model.state_dict(), p['pretext_model']) # Mine the topk nearest neighbors at the very end (Train) # These will be served as input to the SCAN loss. print( colored( 'Fill memory bank for mining the nearest neighbors (train) ...', 'blue')) fill_memory_bank(base_dataloader, model, memory_bank_base) topk = 20 print('Mine the nearest neighbors (Top-%d)' % (topk)) indices, acc = memory_bank_base.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' % (topk, 100 * acc)) np.save(p['topk_neighbors_train_path'], indices) # Mine the topk nearest neighbors at the very end (Val) # These will be used for validation. print( colored('Fill memory bank for mining the nearest neighbors (val) ...', 'blue')) fill_memory_bank(val_dataloader, model, memory_bank_val) topk = 5 print('Mine the nearest neighbors (Top-%d)' % (topk)) indices, acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' % (topk, 100 * acc)) np.save(p['topk_neighbors_val_path'], indices)
print(colored('Retrieve optimizer', 'blue')) optimizer = get_optimizer(p, model) print(optimizer) # Dataset print(colored('Retrieve dataset', 'blue')) # Transforms strong_transforms = get_train_transformations(p) val_transforms = get_val_transformations(p) train_dataset = get_train_dataset(p, {'standard': val_transforms, 'augment': strong_transforms}, split='train', to_augmented_dataset=True) train_dataloader = get_train_dataloader(p, train_dataset) <<<<<<< HEAD #val_dataset = get_val_dataset(p, val_transforms) val_dataloader = get_val_dataloader(p, train_dataset) # print(colored('Train samples %d - Val samples %d' %(len(train_dataset), len(train_dataset)), 'yellow')) #val_ replaced with train_ ======= val_dataset = get_val_dataset(p, val_transforms) val_dataloader = get_val_dataloader(p, val_dataset) print(colored('Train samples %d - Val samples %d' %(len(train_dataset), len(val_dataset)), 'yellow')) >>>>>>> db23360031c529a04f0a144b63e5f3fe49feb44f # Checkpoint if os.path.exists(p['selflabel_checkpoint']): print(colored('Restart from checkpoint {}'.format(p['selflabel_checkpoint']), 'blue')) checkpoint = torch.load(p['selflabel_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch']
def main(): args = FLAGS.parse_args() p = create_config(args.config_env, args.config_exp, args.tb_run) print(colored(p, 'red')) # CUDNN torch.backends.cudnn.benchmark = True # Data print(colored('Get dataset and dataloaders', 'blue')) train_transformations = get_train_transformations(p) val_transformations = get_val_transformations(p) train_dataset = get_train_dataset(p, train_transformations, use_negatives=not p['use_simpred_model'], use_simpred=p['use_simpred_model'], split='train', to_neighbors_dataset=True) val_dataset = get_val_dataset(p, val_transformations, use_negatives=not p['use_simpred_model'], use_simpred=p['use_simpred_model'], to_neighbors_dataset=True) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Train transforms:', train_transformations) print('Validation transforms:', val_transformations) print('Train samples %d - Val samples %d' % (len(train_dataset), len(val_dataset))) # Tensorboard writer writer = SummaryWriter(log_dir=p['scan_tb_dir']) # Model print(colored('Get model', 'blue')) model = get_model(p, p['pretext_model']) print(model) model = torch.nn.DataParallel(model) model = model.cuda() # Simpred Model if p['use_simpred_model']: print(colored('Get simpred model', 'blue')) simpred_model = get_model(p, p['simpred_model'], load_simpred=True) print(simpred_model) simpred_model = torch.nn.DataParallel(simpred_model) simpred_model = simpred_model.cuda() for param in simpred_model.parameters(): param.requires_grad = False else: print('Not using simpred model') simpred_model = None # Optimizer print(colored('Get optimizer', 'blue')) optimizer = get_optimizer(p, model, p['update_cluster_head_only']) print(optimizer) # Warning if p['update_cluster_head_only']: print(colored('WARNING: SCAN will only update the cluster head', 'red')) # Loss function print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() print(criterion) # Checkpoint if os.path.exists(p['scan_checkpoint']): print( colored('Restart from checkpoint {}'.format(p['scan_checkpoint']), 'blue')) checkpoint = torch.load(p['scan_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] best_acc = checkpoint['best_acc'] best_acc_head = checkpoint['best_acc_head'] else: print( colored('No checkpoint file at {}'.format(p['scan_checkpoint']), 'blue')) start_epoch = 0 best_acc = 0 best_acc_head = None # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' % (epoch + 1, p['epochs']), 'yellow')) print(colored('-' * 15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') umcl_train(train_dataloader, model, simpred_model, criterion, optimizer, epoch, writer, p['update_cluster_head_only']) # Evaluate print('Make prediction on validation set ...') predictions = get_predictions(p, val_dataloader, model) print('Evaluate based on similarity accuracy') stats = umcl_evaluate(p, val_dataloader, model, simpred_model) print(stats) highest_acc_head = stats['highest_acc_head'] highest_acc = stats['highest_acc'] if highest_acc > best_acc: print('New highest accuracy on validation set: %.4f -> %.4f' % (best_acc, highest_acc)) print('Highest accuracy head is %d' % highest_acc_head) best_acc = highest_acc best_acc_head = highest_acc_head torch.save( { 'model': model.module.state_dict(), 'head': best_acc_head }, p['scan_model']) else: print('No new highest accuracy on validation set: %.4f -> %.4f' % (best_acc, highest_acc)) print('Highest accuracy head is %d' % highest_acc_head) print('Evaluate with hungarian matching algorithm ...') clustering_stats = hungarian_evaluate(highest_acc_head, predictions, compute_confusion_matrix=False, tf_writer=writer, epoch=epoch) print(clustering_stats) # Checkpoint print('Checkpoint ...') torch.save( { 'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1, 'best_acc': best_acc, 'best_acc_head': best_acc_head }, p['scan_checkpoint']) # Evaluate and save the final model print( colored('Evaluate best model based on similarity accuracy at the end', 'blue')) model_checkpoint = torch.load(p['scan_model'], map_location='cpu') model.module.load_state_dict(model_checkpoint['model']) predictions, features, thumbnails = get_predictions(p, val_dataloader, model, return_features=True, return_thumbnails=True) writer.add_embedding(features, predictions[0]['targets'], thumbnails, p['epochs'], p['scan_tb_dir']) clustering_stats = hungarian_evaluate(model_checkpoint['head'], predictions, class_names=val_dataset.classes, compute_confusion_matrix=True, confusion_matrix_file=os.path.join( p['scan_dir'], 'confusion_matrix.png')) print(clustering_stats)
def main(): # Read config file print(colored('Read config file {} ...'.format(args.config_exp), 'blue')) with open(args.config_exp, 'r') as stream: config = yaml.safe_load(stream) config[ 'batch_size'] = 512 # To make sure we can evaluate on a single 1080ti print(config) # Get dataset print(colored('Get validation dataset ...', 'blue')) transforms = get_val_transformations(config) dataset = get_val_dataset(config, transforms) dataloader = get_val_dataloader(config, dataset) print('Number of samples: {}'.format(len(dataset))) # Get model print(colored('Get model ...', 'blue')) model = get_model(config) print(model) # Read model weights print(colored('Load model weights ...', 'blue')) state_dict = torch.load(args.model, map_location='cpu') if config['setup'] in ['simclr', 'moco', 'selflabel']: model.load_state_dict(state_dict) elif config['setup'] == 'scan': model.load_state_dict(state_dict['model']) else: raise NotImplementedError # CUDA model.cuda() # Perform evaluation if config['setup'] in ['simclr', 'moco']: print( colored( 'Perform evaluation of the pretext task (setup={}).'.format( config['setup']), 'blue')) print('Create Memory Bank') if config['setup'] == 'simclr': # Mine neighbors after MLP memory_bank = MemoryBank(len(dataset), config['model_kwargs']['features_dim'], config['num_classes'], config['criterion_kwargs']['temperature']) else: # Mine neighbors before MLP memory_bank = MemoryBank(len(dataset), config['model_kwargs']['features_dim'], config['num_classes'], config['temperature']) memory_bank.cuda() print('Fill Memory Bank') fill_memory_bank(dataloader, model, memory_bank) print('Mine the nearest neighbors') for topk in [1, 5, 20]: # Similar to Fig 2 in paper _, acc = memory_bank.mine_nearest_neighbors(topk) print( 'Accuracy of top-{} nearest neighbors on validation set is {:.2f}' .format(topk, 100 * acc)) elif config['setup'] in ['scan', 'selflabel']: print( colored( 'Perform evaluation of the clustering model (setup={}).'. format(config['setup']), 'blue')) head = state_dict['head'] if config['setup'] == 'scan' else 0 predictions, features = get_predictions(config, dataloader, model, return_features=True) clustering_stats = hungarian_evaluate(head, predictions, dataset.classes, compute_confusion_matrix=True) print(clustering_stats) if args.visualize_prototypes: prototype_indices = get_prototypes(config, predictions[head], features, model) visualize_indices(prototype_indices, dataset, clustering_stats['hungarian_match']) else: raise NotImplementedError
def main(): args = FLAGS.parse_args() p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # CUDNN torch.backends.cudnn.benchmark = True # Data print(colored('Get dataset and dataloaders', 'blue')) train_transformations = get_train_transformations(p) #val_transformations = get_val_transformations(p) train_dataset = get_train_dataset(p, train_transformations, split='train', to_neighbors_dataset = True) #val_dataset = get_val_dataset(p, val_transformations, to_neighbors_dataset = True) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, train_dataset) #!val_ replaced with train_ print('Train transforms:', train_transformations) #print('Validation transforms:', val_transformations) #print('Train samples %d - Val samples %d' %(len(train_dataset), len(val_dataset))) # Model print(colored('Get model', 'blue')) model = get_model(p, p['pretext_model']) print(model) model = torch.nn.DataParallel(model) model = model.cuda() # Optimizer print(colored('Get optimizer', 'blue')) optimizer = get_optimizer(p, model, p['update_cluster_head_only']) print(optimizer) # Warning if p['update_cluster_head_only']: print(colored('WARNING: SCAN will only update the cluster head', 'red')) # Loss function print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() print(criterion) # Checkpoint if os.path.exists(p['scan_checkpoint']): print(colored('Restart from checkpoint {}'.format(p['scan_checkpoint']), 'blue')) checkpoint = torch.load(p['scan_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] best_loss = checkpoint['best_loss'] best_loss_head = checkpoint['best_loss_head'] else: print(colored('No checkpoint file at {}'.format(p['scan_checkpoint']), 'blue')) start_epoch = 0 best_loss = 1e4 best_loss_head = None # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' %(epoch+1, p['epochs']), 'yellow')) print(colored('-'*15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') scan_train(train_dataloader, model, criterion, optimizer, epoch, p['update_cluster_head_only']) # Evaluate #!!!!!!!!!!!!!!!!!Skipping the next lines because we are not evaluating YET. print('Make prediction on validation set ...') predictions = get_predictions(p, val_dataloader, model) #inputting the train data to get the clusters !!
def main(): args = FLAGS.parse_args() p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # CUDNN torch.backends.cudnn.benchmark = True # Data print(colored('Get dataset and dataloaders', 'blue')) train_transformations = get_train_transformations(p) val_transformations = get_val_transformations(p) print('Train transforms:', train_transformations) print('Validation transforms:', val_transformations) train_dataset = get_train_dataset(p, train_transformations, split='train') val_dataset = get_val_dataset(p, val_transformations) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Train samples %d - Val samples %d' % (len(train_dataset), len(val_dataset))) # Model print(colored('Get model', 'blue')) model = get_model(p) print(model) # Optimizer print(colored('Get optimizer', 'blue')) optimizer = get_optimizer(p, model, p['update_cluster_head_only']) print(optimizer) # Warning if p['update_cluster_head_only']: print( colored( 'WARNING: Linear probing will only update the cluster head', 'red')) # Loss function print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() print(criterion) model = torch.nn.DataParallel(model) model = model.cuda() state = torch.load(p['pretext_model'], map_location='cpu') missing = model.load_state_dict(state, strict=False) print('missing components', missing) if args.mode == 'train': # Checkpoint if os.path.exists(p['linearprobe_checkpoint']): print( colored( 'Restart from checkpoint {}'.format( p['linearprobe_checkpoint']), 'blue')) checkpoint = torch.load(p['linearprobe_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] best_loss = checkpoint['best_loss'] else: print( colored( 'No checkpoint file at {}'.format( p['linearprobe_checkpoint']), 'blue')) start_epoch = 0 best_loss = 1e4 # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' % (epoch + 1, p['epochs']), 'yellow')) print(colored('-' * 15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') linearprobe_train(train_dataloader, model, criterion, optimizer, epoch) if (epoch + 1) % 5 == 0: print('Evaluate based on CE loss ...') linearprobe_stats = linearprobe_evaluate( val_dataloader, model, criterion) loss = linearprobe_stats['loss'] if loss < best_loss: best_loss = loss torch.save({'model': model.module.state_dict()}, p['linearprobe_model']) # Checkpoint print('Checkpoint ...') print(linearprobe_stats) torch.save( { 'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1, 'best_loss': loss }, p['linearprobe_checkpoint']) # Evaluate and save the final model print(colored('Evaluate best model', 'blue')) model_checkpoint = torch.load(p['linearprobe_model'], map_location='cpu') model.module.load_state_dict(model_checkpoint['model']) linearprobe_stats = linearprobe_evaluate(val_dataloader, model, criterion) print(linearprobe_stats) print('Final Accuracy:', linearprobe_stats['accuracy'])
def main(): cv2.setNumThreads(1) # Retrieve config file p = create_config(args.config_env, args.config_exp) sys.stdout = Logger(p['log_file']) print('Python script is {}'.format(os.path.abspath(__file__))) print(colored(p, 'red')) # Get model print(colored('Retrieve model', 'blue')) model = get_model(p) print(model) model = model.cuda() # Freeze all layers except final 1 by 1 convolutional layer for name, param in model.named_parameters(): if name not in ['decoder.4.weight', 'decoder.4.bias']: param.requires_grad = False # Get criterion print(colored('Get loss', 'blue')) criterion = torch.nn.CrossEntropyLoss(ignore_index=255) criterion.cuda() print(criterion) # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Optimizer print(colored('Retrieve optimizer', 'blue')) parameters = list(filter(lambda p: p.requires_grad, model.parameters())) assert len(parameters) == 2 # decoder.4.weight, decoder.4.bias optimizer = get_optimizer(p, parameters) print(optimizer) # Dataset print(colored('Retrieve dataset', 'blue')) train_transforms = get_train_transformations() val_transforms = get_val_transformations() train_dataset = get_train_dataset(p, train_transforms) val_dataset = get_val_dataset(p, val_transforms) true_val_dataset = get_val_dataset( p, None) # True validation dataset without reshape - For validation. train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print( colored( 'Train samples %d - Val samples %d' % (len(train_dataset), len(val_dataset)), 'yellow')) # Resume from checkpoint if os.path.exists(p['checkpoint']): print( colored('Restart from checkpoint {}'.format(p['checkpoint']), 'blue')) checkpoint = torch.load(p['checkpoint'], map_location='cpu') optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['model']) model.cuda() start_epoch = checkpoint['epoch'] best_epoch = checkpoint['best_epoch'] best_iou = checkpoint['best_iou'] else: print( colored('No checkpoint file at {}'.format(p['checkpoint']), 'blue')) start_epoch = 0 best_epoch = 0 best_iou = 0 model = model.cuda() # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' % (epoch + 1, p['epochs']), 'yellow')) print(colored('-' * 10, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') eval_train = train_segmentation_vanilla( p, train_dataloader, model, criterion, optimizer, epoch, freeze_batchnorm=p['freeze_batchnorm']) # Evaluate online -> This will use batched eval where every image is resized to the same resolution. print('Evaluate ...') eval_val = eval_segmentation_supervised_online(p, val_dataloader, model) if eval_val['mIoU'] > best_iou: print('Found new best model: %.2f -> %.2f (mIoU)' % (100 * best_iou, 100 * eval_val['mIoU'])) best_iou = eval_val['mIoU'] best_epoch = epoch torch.save(model.state_dict(), p['best_model']) else: print('No new best model: %.2f -> %.2f (mIoU)' % (100 * best_iou, 100 * eval_val['mIoU'])) print('Last best model was found in epoch %d' % (best_epoch)) # Checkpoint print('Checkpoint ...') torch.save( { 'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1, 'best_epoch': best_epoch, 'best_iou': best_iou }, p['checkpoint']) # Evaluate best model at the end -> This will evaluate the predictions on the original resolution. print(colored('Evaluating best model at the end', 'blue')) model.load_state_dict(torch.load(p['best_model'])) save_results_to_disk(p, val_dataloader, model, crf_postprocess=args.crf_postprocess) eval_stats = eval_segmentation_supervised_offline(p, true_val_dataset, verbose=True)
def main(): # Retrieve config file p = create_config(args) print(colored(p, 'red')) # Model print(colored('Retrieve model', 'blue')) model = get_model(p) print('Model is {}'.format(model.__class__.__name__)) print(model) model = torch.nn.DataParallel(model) model = model.cuda() # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'blue')) transforms = get_val_transformations(p) train_dataset = get_train_dataset(p, transforms) val_dataset = get_val_dataset(p, transforms) train_dataloader = get_val_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Dataset contains {}/{} train/val samples'.format( len(train_dataset), len(val_dataset))) # Memory Bank print(colored('Build MemoryBank', 'blue')) memory_bank_train = MemoryBank(len(train_dataset), 2048, p['num_classes'], p['temperature']) memory_bank_train.cuda() memory_bank_val = MemoryBank(len(val_dataset), 2048, p['num_classes'], p['temperature']) memory_bank_val.cuda() # Load the official MoCoV2 checkpoint print(colored('Downloading moco v2 checkpoint', 'blue')) os.system( 'wget -L https://dl.fbaipublicfiles.com/moco/moco_checkpoints/moco_v2_800ep/moco_v2_800ep_pretrain.pth.tar' ) moco_state = torch.load('moco_v2_800ep_pretrain.pth.tar', map_location='cpu') # Transfer moco weights print(colored('Transfer MoCo weights to model', 'blue')) new_state_dict = {} state_dict = moco_state['state_dict'] for k in list(state_dict.keys()): # Copy backbone weights if k.startswith('module.encoder_q' ) and not k.startswith('module.encoder_q.fc'): new_k = 'module.backbone.' + k[len('module.encoder_q.'):] new_state_dict[new_k] = state_dict[k] # Copy mlp weights elif k.startswith('module.encoder_q.fc'): new_k = 'module.contrastive_head.' + k[len('module.encoder_q.fc.' ):] new_state_dict[new_k] = state_dict[k] else: raise ValueError('Unexpected key {}'.format(k)) model.load_state_dict(new_state_dict) os.system('rm -rf moco_v2_800ep_pretrain.pth.tar') # Save final model print(colored('Save pretext model', 'blue')) torch.save(model.module.state_dict(), p['pretext_model']) model.module.contrastive_head = torch.nn.Identity( ) # In this case, we mine the neighbors before the MLP. # Mine the topk nearest neighbors (Train) # These will be used for training with the SCAN-Loss. topk = 50 print( colored('Mine the nearest neighbors (Train)(Top-%d)' % (topk), 'blue')) transforms = get_val_transformations(p) train_dataset = get_train_dataset(p, transforms) fill_memory_bank(train_dataloader, model, memory_bank_train) indices, acc = memory_bank_train.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' % (topk, 100 * acc)) np.save(p['topk_neighbors_train_path'], indices) # Mine the topk nearest neighbors (Validation) # These will be used for validation. topk = 5 print(colored('Mine the nearest neighbors (Val)(Top-%d)' % (topk), 'blue')) fill_memory_bank(val_dataloader, model, memory_bank_val) print('Mine the neighbors') indices, acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' % (topk, 100 * acc)) np.save(p['topk_neighbors_val_path'], indices)
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # Model print(colored('Retrieve model', 'green')) model = get_model(p) print('Model is {}'.format(model.__class__.__name__)) print(model) # model = torch.nn.DataParallel(model) model = model.to(device) # CUDNN print(colored('Set CuDNN benchmark', 'green')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'green')) transforms = get_val_transformations(p) train_dataset = get_train_dataset(p, transforms) val_dataset = get_val_dataset(p, transforms) train_dataloader = get_val_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Dataset contains {}/{} train/val samples'.format(len(train_dataset), len(val_dataset))) # Memory Bank print(colored('Build MemoryBank', 'green')) memory_bank_train = MemoryBank(len(train_dataset), 2048, p['num_classes'], p['temperature']) memory_bank_train.to(device) memory_bank_val = MemoryBank(len(val_dataset), 2048, p['num_classes'], p['temperature']) memory_bank_val.to(device) # Load the official MoCoV2 checkpoint print(colored('Downloading moco v2 checkpoint', 'green')) # os.system('wget -L https://dl.fbaipublicfiles.com/moco/moco_checkpoints/moco_v2_800ep/moco_v2_800ep_pretrain.pth.tar') # Uploaded the model to Mist : Johan moco_state = torch.load(main_dir + model_dir + 'moco_v2_800ep_pretrain.pth.tar', map_location=device) # Transfer moco weights print(colored('Transfer MoCo weights to model', 'green')) new_state_dict = {} state_dict = moco_state['state_dict'] # for k in list(state_dict.keys()): # # Copy backbone weights # if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'): # new_k = 'module.backbone.' + k[len('module.encoder_q.'):] # new_state_dict[new_k] = state_dict[k] # # # Copy mlp weights # elif k.startswith('module.encoder_q.fc'): # new_k = 'module.contrastive_head.' + k[len('module.encoder_q.fc.'):] # new_state_dict[new_k] = state_dict[k] # # else: # raise ValueError('Unexpected key {}'.format(k)) #Changed by Johan for k, v in state_dict.items(): if "conv" in k or "bn" in k or "layer" in k: new_k = "backbone." + k.split("module.encoder_q.")[1] new_state_dict[new_k] = v else: new_k = "contrastive_head." + k.split("module.encoder_q.fc.")[1] new_state_dict[new_k] = v model.load_state_dict(new_state_dict) # os.system('rm -rf moco_v2_800ep_pretrain.pth.tar') # Save final model print(colored('Save pretext model', 'green')) torch.save(model.state_dict(), p['pretext_model']) # model.contrastive_head = torch.nn.Identity() # In this case, we mine the neighbors before the MLP. model.contrastive_head = Identity()
def main(): args = FLAGS.parse_args() p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # CUDNN torch.backends.cudnn.benchmark = True # Data print(colored('Get dataset and dataloaders', 'blue')) train_transformations = get_train_transformations(p) val_transformations = get_val_transformations(p) train_dataset = get_train_dataset(p, train_transformations, split='train', to_neighbors_strangers_dataset = True) val_dataset = get_val_dataset(p, val_transformations, to_neighbors_strangers_dataset = True) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Train transforms:', train_transformations) print('Validation transforms:', val_transformations) print('Train samples %d - Val samples %d' %(len(train_dataset), len(val_dataset))) # Model print(colored('Get model', 'blue')) model = get_model(p, p['pretext_model']) print(model) model = torch.nn.DataParallel(model) model = model.cuda() # Optimizer print(colored('Get optimizer', 'blue')) optimizer = get_optimizer(p, model, p['update_cluster_head_only']) print(optimizer) # Warning if p['update_cluster_head_only']: print(colored('WARNING: SCAN will only update the cluster head', 'red')) # Loss function print(colored('Get loss', 'blue')) criterion = get_criterion(p) criterion.cuda() print(criterion) if args.mode == 'train': # Checkpoint if os.path.exists(p['scanf_checkpoint']): print(colored('Restart from checkpoint {}'.format(p['scanf_checkpoint']), 'blue')) checkpoint = torch.load(p['scanf_checkpoint'], map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] best_loss = checkpoint['best_loss'] best_loss_head = checkpoint['best_loss_head'] else: print(colored('No checkpoint file at {}'.format(p['scanf_checkpoint']), 'blue')) start_epoch = 0 best_loss = 1e4 best_loss_head = None # Main loop print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' %(epoch+1, p['epochs']), 'yellow')) print(colored('-'*15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') scanf_train(train_dataloader, model, criterion, optimizer, epoch, p['update_cluster_head_only']) # Evaluate print('Make prediction on validation set ...') predictions = get_predictions(p, val_dataloader, model) print('Evaluate based on SCAN loss ...') scanf_stats = scanf_evaluate(predictions) print(scanf_stats) lowest_loss_head = scanf_stats['lowest_loss_head'] lowest_loss = scanf_stats['lowest_loss'] if lowest_loss < best_loss: print('New lowest loss on validation set: %.4f -> %.4f' %(best_loss, lowest_loss)) print('Lowest loss head is %d' %(lowest_loss_head)) best_loss = lowest_loss best_loss_head = lowest_loss_head torch.save({'model': model.module.state_dict(), 'head': best_loss_head}, p['scanf_model']) else: print('No new lowest loss on validation set: %.4f -> %.4f' %(best_loss, lowest_loss)) print('Lowest loss head is %d' %(best_loss_head)) print('Evaluate with hungarian matching algorithm ...') clustering_stats = hungarian_evaluate(lowest_loss_head, predictions, compute_confusion_matrix=False) print(clustering_stats) # Checkpoint print('Checkpoint ...') torch.save({'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1, 'best_loss': best_loss, 'best_loss_head': best_loss_head}, p['scanf_checkpoint']) # Evaluate and save the final model print(colored('Evaluate best model based on SCAN metric at the end', 'blue')) model_checkpoint = torch.load(p['scanf_model'], map_location='cpu') model.module.load_state_dict(model_checkpoint['model']) predictions = get_predictions(p, val_dataloader, model) gt_targets = predictions[model_checkpoint['head']]['targets'] cluster_predictions = predictions[model_checkpoint['head']]['predictions'] print(gt_targets.shape) print(cluster_predictions.shape) torch.save(gt_targets, 'scanf_gt_targets.pth.tar') torch.save(cluster_predictions, 'scanf_cluster_predictions.pth.tar') clustering_stats = hungarian_evaluate(model_checkpoint['head'], predictions, class_names=val_dataset.dataset.classes, compute_confusion_matrix=True, confusion_matrix_file=os.path.join(p['scanf_dir'], 'confusion_matrix.png')) print(clustering_stats) print('Final Accuracy:', clustering_stats['ACC'])
def main(args): # Retrieve config file p = create_config(args.config_env, args.config_exp) print(colored(p, 'red')) # Model print(colored('Retrieve model', 'blue')) model = get_model(p) print('Model is {}'.format(model.__class__.__name__)) print('Model parameters: {:.2f}M'.format(sum(p.numel() for p in model.parameters()) / 1e6)) print(model) # from torchsummary import summary # summary(model, (3, p['transformation_kwargs']['crop_size'], p['transformation_kwargs']['crop_size'])) model = model.cuda() # CUDNN print(colored('Set CuDNN benchmark', 'blue')) torch.backends.cudnn.benchmark = True # Dataset print(colored('Retrieve dataset', 'blue')) train_transforms = get_train_transformations(p) print('Train transforms:', train_transforms) val_transforms = get_val_transformations(p) print('Validation transforms:', val_transforms) train_dataset = get_train_dataset(p, train_transforms, to_augmented_dataset=True, split='train+unlabeled') # Split is for stl-10 val_dataset = get_val_dataset(p, val_transforms) train_dataloader = get_train_dataloader(p, train_dataset) val_dataloader = get_val_dataloader(p, val_dataset) print('Dataset contains {}/{} train/val samples'.format(len(train_dataset), len(val_dataset))) # Criterion print(colored('Retrieve criterion', 'blue')) criterion = get_criterion(p) print('Criterion is {}'.format(criterion.__class__.__name__)) criterion = criterion.cuda() # Checkpoint # p['pretext_checkpoint'] = p['pretext_checkpoint'].replace('checkpoint.pth.tar', '2nd_94306c9_checkpoint.pth.tar') # Specific model assert os.path.exists(p['pretext_checkpoint']), "Checkpoint not found - can't fine-tune." print(colored('Restart from checkpoint {}'.format(p['pretext_checkpoint']), 'blue')) checkpoint = torch.load(p['pretext_checkpoint'], map_location='cpu') # optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['model']) model.cuda() # start_epoch = checkpoint['epoch'] start_epoch = 0 # Train linear model from representations to evaluate attributes classification print(colored('Train linear', 'blue')) for parameter in model.parameters(): parameter.requires_grad = False # model = nn.Sequential(model, AttributesHead(p['model_kwargs']['features_dim'], p['num_attribute_classes'])) model.contrastive_head = nn.Sequential(model.contrastive_head, nn.Linear(p['model_kwargs']['features_dim'], p['num_attribute_classes'])) model.cuda() # Optimizer and scheduler print(colored('Retrieve optimizer', 'blue')) optimizer = get_optimizer(p, model) print(optimizer) # Training print(colored('Starting main loop', 'blue')) for epoch in range(start_epoch, p['epochs']): print(colored('Epoch %d/%d' % (epoch, p['epochs']), 'yellow')) print(colored('-' * 15, 'yellow')) # Adjust lr lr = adjust_learning_rate(p, optimizer, epoch) print('Adjusted learning rate to {:.5f}'.format(lr)) # Train print('Train ...') simclr_fine_tune_train(train_dataloader, model, criterion, optimizer, epoch) # Evaluate acc = attributes_evaluate(val_dataloader, model) print('Val set accuracy %.2f' % acc) # Checkpoint print('Checkpoint ...') torch.save({'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 'epoch': epoch + 1}, p['pretext_fine_tune_checkpoint']) # Save final model torch.save(model.state_dict(), p['pretext_fine_tune_model'])