if __name__ == "__main__":
if args.model == 'MnistCNN':
model = MnistCNN()
train_transform, test_transform = get_data_transform('mnist')
train_dataset = datasets.MNIST(args.data_dir,
train=True,
download=False,
transform=train_transform)
test_dataset = datasets.MNIST(args.data_dir,
train=False,
download=False,
transform=test_transform)
elif args.model == 'AlexNet':
model = AlexNetForCIFAR()
train_transform, test_transform = get_data_transform('cifar')
train_dataset = datasets.CIFAR10(args.data_dir,
train=True,
download=False,
transform=train_transform)
test_dataset = datasets.CIFAR10(args.data_dir,
train=False,
download=False,
transform=test_transform)
elif args.model == 'LROnMnist':
model = ResNetOnCifar10.LROnMnist()
train_transform, test_transform = get_data_transform('mnist')
train_dataset = datasets.MNIST(args.data_dir,
train=True,
test_dataset = datasets.MNIST(args.data_dir, train=False, download=False,
transform=test_transform)
elif args.model == 'LROnCifar10':
model = ResNetOnCifar10.LROnCifar10()
train_transform, test_transform = get_data_transform('cifar')
train_dataset = datasets.CIFAR10(args.data_dir, train=True, download=False,
transform=train_transform)
test_dataset = datasets.CIFAR10(args.data_dir, train=False, download=False,
transform=test_transform)
elif args.model == 'AlexNet':
train_transform, test_transform = get_data_transform('cifar')
if args.data_name == 'cifar10':
model = AlexNetForCIFAR()
train_dataset = datasets.CIFAR10(args.data_dir, train=True, download=False,
transform=train_transform)
test_dataset = datasets.CIFAR10(args.data_dir, train=False, download=False,
transform=test_transform)
else:
model = AlexNetForCIFAR(num_classes=100)
train_dataset = datasets.CIFAR100(args.data_dir, train=True, download=False,
transform=train_transform)
test_dataset = datasets.CIFAR100(args.data_dir, train=False, download=False,
transform=test_transform)
elif args.model == 'ResNet18OnCifar10':
model = ResNetOnCifar10.ResNet18()
train_transform, test_transform = get_data_transform('cifar')
train_dataset = datasets.CIFAR10(args.data_dir, train=True, download=False,
