def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp, args.tb_run) 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']) # 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)) knn_indices, knn_acc = memory_bank_base.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' % (topk, 100 * knn_acc)) np.save(p['topk_neighbors_train_path'], knn_indices) if p['compute_negatives']: topk = 200 kfn_indices, kfn_acc = memory_bank_base.mine_negatives(topk) print('Accuracy of top-%d furthest neighbors on train set is %.2f' % (topk, 100 * kfn_acc)) np.save(p['topk_furthest_train_path'], kfn_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)) knn_indices, knn_acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' % (topk, 100 * knn_acc)) np.save(p['topk_neighbors_val_path'], knn_indices) if p['compute_negatives']: kfn_indices, kfn_acc = memory_bank_val.mine_negatives(topk) print('Accuracy of top-%d furthest neighbors on val set is %.2f' % (topk, 100 * kfn_acc)) np.save(p['topk_furthest_val_path'], kfn_indices)
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp, args.tb_run) 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')) moco_state = torch.load(p['pretrained'], 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: pass # just silently discard unexpected keys # raise ValueError('Unexpected key {}'.format(k)) model.load_state_dict(new_state_dict) # 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) knn_indices, knn_acc = memory_bank_train.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' % (topk, 100 * knn_acc)) np.save(p['topk_neighbors_train_path'], knn_indices) if p['compute_negatives']: topk = 350 kfn_indices, kfn_acc = memory_bank_train.mine_negatives(topk) print('Accuracy of top-%d furthest neighbors on train set is %.2f' % (topk, 100 * kfn_acc)) np.save(p['topk_furthest_train_path'], kfn_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') knn_indices, knn_acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' % (topk, 100 * knn_acc)) np.save(p['topk_neighbors_val_path'], knn_indices) if p['compute_negatives']: kfn_indices, kfn_acc = memory_bank_val.mine_negatives(topk) print('Accuracy of top-%d furthest neighbors on val set is %.2f' % (topk, 100 * kfn_acc)) np.save(p['topk_furthest_val_path'], kfn_indices)
def main(): # Retrieve config file p = create_config(args.config_env, args.config_exp, args.tb_run) 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)) knn_indices, knn_acc = memory_bank_base.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on train set is %.2f' % (topk, 100 * knn_acc)) np.save(p['topk_neighbors_train_path'], knn_indices) if p['compute_negatives']: topk = 200 kfn_indices, kfn_acc = memory_bank_base.mine_negatives(topk) print('Accuracy of top-%d furthest neighbors on train set is %.2f' % (topk, 100 * kfn_acc)) np.save(p['topk_furthest_train_path'], kfn_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') knn_indices, knn_acc = memory_bank_val.mine_nearest_neighbors(topk) print('Accuracy of top-%d nearest neighbors on val set is %.2f' % (topk, 100 * knn_acc)) np.save(p['topk_neighbors_val_path'], knn_indices) if p['compute_negatives']: kfn_indices, kfn_acc = memory_bank_val.mine_negatives(topk) print('Accuracy of top-%d furthest neighbors on val set is %.2f' % (topk, 100 * kfn_acc)) np.save(p['topk_furthest_val_path'], kfn_indices)