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)