def finetune(classifier, dataset, device, args):
params = []
if args.finetune_mode == 'freeze':
print('[INFO] Finetune classifier only for the last layer...')
for name, param in classifier.named_parameters():
if 'encoder' in name or 'agg' in name:
param.requires_grad = False
else:
params.append({'params': param})
elif args.finetune_mode == 'smaller':
print('[INFO] Finetune the whole classifier where the backbone have a smaller lr...')
for name, param in classifier.named_parameters():
if 'encoder' in name or 'agg' in name:
params.append({'params': param, 'lr': args.lr / 10})
else:
params.append({'params': param})
else:
print('[INFO] Finetune the whole classifier...')
for name, param in classifier.named_parameters():
params.append({'params': param})
if args.optimizer == 'sgd':
optimizer = optim.SGD(params, lr=args.lr, weight_decay=args.wd, momentum=args.momentum)
elif args.optimizer == 'adam':
optimizer = optim.Adam(params, lr=args.lr, weight_decay=args.wd, betas=(0.9, 0.98), eps=1e-09,
amsgrad=True)
else:
raise ValueError('Invalid optimizer!')
criterion = nn.CrossEntropyLoss().cuda(device)
sampled_indices = np.arange(len(dataset))
np.random.shuffle(sampled_indices)
sampled_indices = sampled_indices[:int(len(sampled_indices) * args.finetune_ratio)]
data_loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers,
shuffle=False, pin_memory=True, drop_last=True,
sampler=SubsetRandomSampler(sampled_indices))
classifier.train()
for epoch in range(args.finetune_epochs):
losses = []
accuracies = []
with tqdm(data_loader, desc=f'EPOCH [{epoch + 1}/{args.finetune_epochs}]') as progress_bar:
for x, y in progress_bar:
x, y = x.cuda(device, non_blocking=True), y.cuda(device, non_blocking=True)
x = x.view(x.shape[0] * x.shape[1], *x.shape[2:])
out = classifier(x)
loss = criterion(out, y.view(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
accuracies.append(
logits_accuracy(out, y.view(-1), topk=(1,))[0])
progress_bar.set_postfix({'Loss': np.mean(losses), 'Acc': np.mean(accuracies)})
def pretrain(run_id, model, dataset, device, args):
if args.optimizer == 'sgd':
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
weight_decay=args.wd,
momentum=args.momentum)
elif args.optimizer == 'adam':
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.wd,
betas=(0.9, 0.98),
eps=1e-09,
amsgrad=True)
else:
raise ValueError('Invalid optimizer!')
criterion = nn.CrossEntropyLoss().cuda(device)
data_loader = DataLoader(dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True,
drop_last=True)
model.train()
for epoch in range(args.pretrain_epochs):
losses = []
accuracies = []
adjust_learning_rate(optimizer, args.lr, epoch, args.pretrain_epochs,
args)
with tqdm(data_loader,
desc=f'EPOCH [{epoch + 1}/{args.pretrain_epochs}]'
) as progress_bar:
for x, _ in progress_bar:
x = x.cuda(device, non_blocking=True)
output, target = model(x)
loss = criterion(output, target)
acc = logits_accuracy(output, target, topk=(1, ))[0]
accuracies.append(acc)
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
progress_bar.set_postfix({
'Loss': np.mean(losses),
'Acc': np.mean(accuracies)
})
