if args.test_only or len(args.resume) > 0:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
checkpoint = torch.load('./checkpoint/' + args.resume)
net.load_state_dict(checkpoint['net'])
lemniscate = checkpoint['lemniscate']
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch']
# define loss function
criterion = NCECriterion()
net.to(device)
lemniscate.to(device)
criterion.to(device)
if args.test_only:
acc = kNN(0, net, lemniscate, trainloader, testloader, 200, args.nce_t, 1)
sys.exit(0)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
def adjust_learning_rate(optimizer, epoch):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
lr = args.lr
if epoch >= 20:
def build_model():
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
if args.architecture == 'resnet18':
net = models.__dict__['resnet18_cifar'](low_dim=args.low_dim)
elif args.architecture == 'wrn-28-2':
net = models.WideResNet(depth=28,
num_classes=args.low_dim,
widen_factor=2,
dropRate=0).to(args.device)
elif args.architecture == 'wrn-28-10':
net = models.WideResNet(depth=28,
num_classes=args.low_dim,
widen_factor=10,
dropRate=0).to(args.device)
# define leminiscate
if args.nce_k > 0:
lemniscate = NCEAverage(args.low_dim, args.ndata, args.nce_k,
args.nce_t, args.nce_m)
else:
lemniscate = LinearAverage(args.low_dim, args.ndata, args.nce_t,
args.nce_m)
if args.device == 'cuda':
net = torch.nn.DataParallel(net,
device_ids=range(
torch.cuda.device_count()))
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay,
nesterov=True)
# Model
if args.test_only or len(args.resume) > 0:
# Load checkpoint.
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.resume)
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
lemniscate = checkpoint['lemniscate']
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch'] + 1
if args.lr_scheduler == 'multi-step':
if args.epochs == 200:
steps = [60, 120, 160]
elif args.epochs == 600:
steps = [180, 360, 480, 560]
else:
raise RuntimeError(
f"need to config steps for epoch = {args.epochs} first.")
scheduler = lr_scheduler.MultiStepLR(optimizer,
steps,
gamma=0.2,
last_epoch=start_epoch - 1)
elif args.lr_scheduler == 'cosine':
scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
args.epochs,
eta_min=0.00001,
last_epoch=start_epoch - 1)
elif args.lr_scheduler == 'cosine-with-restart':
scheduler = CosineAnnealingLRWithRestart(optimizer,
eta_min=0.00001,
last_epoch=start_epoch - 1)
else:
raise ValueError("not supported")
# define loss function
if hasattr(lemniscate, 'K'):
criterion = NCECriterion(args.ndata)
else:
criterion = nn.CrossEntropyLoss()
net.to(args.device)
lemniscate.to(args.device)
criterion.to(args.device)
return net, lemniscate, optimizer, criterion, scheduler, best_acc, start_epoch
