def set_model(opt):
model = SupCEResNet(name=opt.model, num_classes=opt.n_cls)
criterion = torch.nn.CrossEntropyLoss()
if opt.eval:
ckpt = torch.load(opt.ckpt, map_location='cpu')
state_dict = ckpt['model']
state_dict = {
k.replace("module.", ""): v
for k, v in state_dict.items()
}
model.load_state_dict(state_dict, strict=True)
# enable synchronized Batch Normalization
if opt.syncBN:
model = apex.parallel.convert_syncbn_model(model)
if torch.cuda.is_available():
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model = model.cuda()
criterion = criterion.cuda()
cudnn.benchmark = True
return model, criterion
def set_model(opt):
model = SupCEResNet(name=opt.model, num_classes=opt.num_classes)
if torch.cuda.is_available():
if torch.cuda.device_count() > 1:
print('available device:', torch.cuda.device_count())
model = torch.nn.DataParallel(model)
model = model.cuda()
cudnn.benchmark = True
return model
def set_model(opt):
model = SupCEResNet(name=opt.model, num_classes=opt.n_cls)
criterion = torch.nn.CrossEntropyLoss()
# enable synchronized Batch Normalization
if opt.syncBN:
model = apex.parallel.convert_syncbn_model(model)
if torch.cuda.is_available():
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model = model.cuda()
criterion = criterion.cuda()
cudnn.benchmark = True
return model, criterion
def set_model(opt):
if (opt.model == 'WRN28'):
model = SupCEWRN(name=opt.model, num_classes=opt.n_cls)
else:
model = SupCEResNet(name=opt.model, num_classes=opt.n_cls)
criterion = GraphLoss(opt.temp)
# enable synchronized Batch Normalization
if opt.syncBN:
model = apex.parallel.convert_syncbn_model(model)
if torch.cuda.is_available():
#if torch.cuda.device_count() > 1:
# print('available device:', torch.cuda.device_count())
# model = torch.nn.DataParallel(model)
model = model.cuda()
criterion = criterion.cuda()
cudnn.benchmark = True
state_dict = torch.load(opt.model_name + '/best.pth')['model']
#state_dict = torch.load('./resnet18_clean_acc93.pth')['model']
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
if (k[0:7] == 'module.'):
name = k[7:] # remove `module.`
new_state_dict[name] = v
else:
name = k
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
return model, criterion