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) })