def load_model():
if args.type == 'OE':
model = resnet_oe.ResNet18(dataset='CIFAR10').gpu()
else:
model = resnet.ResNet18().cuda()
model.load_state_dict(
torch.load(f'./pretrained_models/CIFAR10_{args.type}.pt'))
model.eval()
return model
def get_dkl_model(dataset='MNIST', binary=False):
num_classes = 2 if binary else (100 if dataset == 'CIFAR100' else 10)
feature_extractor = LeNetMadry(
binary=False,
feature_extractor=True) if dataset == 'MNIST' else resnet.ResNet18(
num_classes=num_classes, feature_extractor=True)
feature_extractor.cuda()
feature_extractor.train()
num_features = feature_extractor.fc2.in_features if dataset == 'MNIST' else feature_extractor.linear.in_features
model = DKLModel(feature_extractor, num_dim=num_features)
likelihood = gpytorch.likelihoods.SoftmaxLikelihood(
num_features=model.num_dim, num_classes=num_classes)
return model, likelihood
def load_model():
model = resnet.ResNet18(num_classes=2).cuda()
model.load_state_dict(torch.load(f'./pretrained_models/binary_CIFAR10.pt'))
model.eval()
return model
class1, class2 = 3, 9
train_loader = dl.binary_SVHN(class1, class2, train=True)
val_loader, test_loader = dl.binary_SVHN(class1, class2, train=False, augm_flag=False, val_size=1000)
elif args.dataset == 'CIFAR100':
print('a')
class1, class2 = 47, 52
train_loader = dl.binary_CIFAR100(class1, class2, train=True)
val_loader, test_loader = dl.binary_CIFAR100(class1, class2, train=False, augm_flag=False, val_size=100)
print(len(train_loader.dataset), len(val_loader.dataset), len(test_loader.dataset))
targets = torch.cat([y for x, y in test_loader], dim=0).numpy()
targets_val = torch.cat([y for x, y in val_loader], dim=0).numpy()
print(len(train_loader.dataset), len(val_loader.dataset), len(test_loader.dataset))
model = LeNetMadry(binary=True) if args.dataset == 'MNIST' else resnet.ResNet18(num_classes=2)
model.cuda()
model.train()
if args.dataset == 'MNIST':
opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
else:
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
criterion = torch.nn.BCEWithLogitsLoss()
pbar = trange(100)
for epoch in pbar:
if epoch+1 in [50,75,90]:
for group in opt.param_groups:
group['lr'] *= .1