def set_model(args, n_data):
# set the model
if args.model == 'c3d':
model = C3D(with_classifier=False)
elif args.model == 'r3d':
model = R3DNet(layer_sizes=(1, 1, 1, 1), with_classifier=False)
elif args.model == 'r21d':
model = R2Plus1DNet(layer_sizes=(1, 1, 1, 1), with_classifier=False)
if args.intra_neg:
contrast = NCEAverage(args.feat_dim, n_data, args.nce_k, args.nce_t,
args.nce_m, args.softmax)
else:
contrast = NCEAverage_ori(args.feat_dim, n_data, args.nce_k,
args.nce_t, args.nce_m, args.softmax)
criterion_1 = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)
criterion_2 = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)
# GPU mode
model = model.cuda()
contrast = contrast.cuda()
criterion_1 = criterion_1.cuda()
criterion_2 = criterion_2.cuda()
cudnn.benchmark = True
return model, contrast, criterion_1, criterion_2
print('==> Building model..')
net = models.__dict__['ResNet18'](low_dim=args.low_dim)
# define leminiscate
lemniscate = NCEAverage(args.low_dim,
ndata,
args.nce_t,
args.nce_m,
batchSize=args.b)
# define loss function
criterion = NCECriterion()
if device == 'cuda':
net = torch.nn.DataParallel(net.cuda())
lemniscate.cuda()
criterion.cuda()
cudnn.benchmark = True
# Model
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']
if args.test_only:
# define leminiscate
if args.nce_k > 0:
lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m)
else:
lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m)
# define loss function
if hasattr(lemniscate, 'K'):
criterion = NCECriterion(ndata)
else:
criterion = nn.CrossEntropyLoss()
if use_cuda:
net.cuda()
net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
lemniscate.cuda()
criterion.cuda()
cudnn.benchmark = True
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 >= 80:
lr = args.lr * (0.1 ** ((epoch-80) // 40))
print(lr)