def set_model(args):
if args.model.startswith('alexnet'):
model = MyAlexNetCMC()
classifier = LinearClassifierAlexNet(layer=args.layer, n_label=args.n_label, pool_type='max')
elif args.model.startswith('resnet'):
model = MyResNetsCMC(name=args.model, view=args.view)
if args.model.endswith('v1'):
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 1)
elif args.model.endswith('v2'):
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 2)
elif args.model.endswith('v3'):
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 4)
elif 'ttt' in args.model:
classifier = LinearClassifierResNet(10, args.n_label, 'avg', 1)
else:
raise NotImplementedError('model not supported {}'.format(args.model))
else:
raise NotImplementedError('model not supported {}'.format(args.model))
# load pre-trained model
print('==> loading pre-trained model')
ckpt = torch.load(args.model_path)
model.load_state_dict(ckpt['model'])
print("==> loaded checkpoint '{}' (epoch {})".format(args.model_path, ckpt['epoch']))
print('==> done')
model = model.cuda()
classifier = classifier.cuda()
model.eval()
return model, classifier
def set_model(args, n_data):
# set the model
if args.model == 'alexnet':
model = MyAlexNetCMC(args.feat_dim)
elif args.model.startswith('resnet'):
model = MyResNetsCMC(args.model)
else:
raise ValueError('model not supported yet {}'.format(args.model))
contrast = NCEAverage(args.feat_dim, n_data, args.nce_k, args.nce_t, args.nce_m, args.softmax)
criterion_l = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)
criterion_ab = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)
if torch.cuda.is_available():
model = model.cuda()
contrast = contrast.cuda()
criterion_ab = criterion_ab.cuda()
criterion_l = criterion_l.cuda()
cudnn.benchmark = True
return model, contrast, criterion_ab, criterion_l