def set_model(args):
model = WideResnet(n_classes=10 if args.dataset == 'CIFAR10' else 100,
k=args.wresnet_k,
n=args.wresnet_n) # wresnet-28-2
name = 'simclr_trained_good_h2.pt'
model.load_state_dict(torch.load(name))
print('model loaded')
model.train()
model.cuda()
criteria_x = nn.CrossEntropyLoss().cuda()
criteria_u = nn.CrossEntropyLoss(reduction='none').cuda()
criteria_z = NT_Xent(args.batchsize, args.temperature, args.mu)
return model, criteria_x, criteria_u, criteria_z
def set_model(args):
model = WideResnet(n_classes=10 if args.dataset == 'CIFAR10' else 100,
k=args.wresnet_k, n=args.wresnet_n) # wresnet-28-2
model.train()
model.cuda()
criteria_x = nn.CrossEntropyLoss().cuda()
criteria_u = nn.CrossEntropyLoss(reduction='none').cuda()
return model, criteria_x, criteria_u