}
os.makedirs(DATASET_DIR['train'], exist_ok=True)
os.makedirs(DATASET_DIR['val'], exist_ok=True)
# Load the data
BANDWIDTH = 16
trainloader, testloader, trainset, testset, _, _, proc_mean, proc_std = get_processed_dataset_loaders(
lambda x: dct_low_pass(x, bandwidth=BANDWIDTH), DATASET, DATASET_DIR,
BATCH_SIZE)
####################
# Select a Network #
####################
# Normalization layer
trans = TransformLayer(mean=proc_mean, std=proc_std)
# Load a model
model = ResNet18(
) # check inside the model_class.cifar10 package for other network options
# If pretrained
if PRETRAINED:
print('---> Working on a pretrained network')
model.load_state_dict(torch.load(PRETRAINED_PATH, map_location='cpu'))
model = model.to(DEVICE)
model.eval()
# If not pretrained, then train it
if not PRETRAINED:
# Specify the path of the dataset. For MNIST and CIFAR-10 the train and validation paths can be the same.
# For ImageNet, please specify to proper train and validation paths.
DATASET_DIR = {'train': 'path-to-ILSVRC2012-train-data',
'val': 'path-to-ILSVRC2012-val-data'
}
# Load the data
trainloader, testloader, trainset, testset, mean, std = get_dataset_loaders(DATASET, DATASET_DIR, BATCH_SIZE)
####################
# Select a Network #
####################
# Normalization layer
trans = TransformLayer(mean=mean, std=std)
# Load a pretrained model
model = resnet50()
model.load_state_dict(torch.load(PRETRAINED_PATH, map_location='cpu'))
model = model.to(DEVICE)
model.eval()
##################################
# Compute margin along subspaces #
##################################
# Create a list of subspaces to evaluate the margin on
SUBSPACE_DIM = 16
DIM = 224