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():
# 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', '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_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', '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)
# Criterion
print(colored('Retrieve criterion', 'green'))
criterion = get_criterion(p)
print('Criterion is {}'.format(criterion.__class__.__name__))
criterion = criterion.to(device)
# Optimizer and scheduler
print(colored('Retrieve optimizer', 'green'))
optimizer = get_optimizer(p, model)
print(optimizer)
# 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'))
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) ...',
'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)
