architectures.py

'''
- this file is necessary to get the right architectures for the models
- it is based on the publicly available code https://github.com/locuslab/smoothing/blob/master/code/architectures.py written by Jeremy Cohen
'''

import torch
from torchvision.models.resnet import resnet50
import torch.backends.cudnn as cudnn
from archs.cifar_resnet import resnet as resnet_cifar
from datasets import get_normalize_layer
from datasets import get_input_center_layer # added for certain SmoothAdv models 
from torch.nn.functional import interpolate

# resnet50 - the classic ResNet-50, sized for ImageNet
# cifar_resnet20 - a 20-layer residual network sized for CIFAR
# cifar_resnet110 - a 110-layer residual network sized for CIFAR
ARCHITECTURES = ["resnet50", "cifar_resnet20", "cifar_resnet110", "cifar_resnet110_selection", "cifar_resnet110_binary"]

def get_architecture(arch: str, dataset: str) -> torch.nn.Module:
    """ Return a neural network (with random weights)
    :param arch: the architecture - should be in the ARCHITECTURES list above
    :param dataset: the dataset - should be in the datasets.DATASETS list
    :return: a Pytorch module
    """
    if arch == "resnet50" and dataset == "imagenet":
        model = torch.nn.DataParallel(resnet50(pretrained=False)).cuda()
        cudnn.benchmark = True
    elif arch == "cifar_resnet20":
        model = resnet_cifar(depth=20, num_classes=10).cuda()
    elif arch == "cifar_resnet110":
        model = resnet_cifar(depth=110, num_classes=10).cuda()
    elif arch == "cifar_resnet110_binary":
        model = resnet_cifar(depth=110, num_classes=2).cuda()
    elif arch == "cifar_resnet110_selection":
        model = resnet_cifar(depth=110, num_classes=1).cuda()
    normalize_layer = get_normalize_layer(dataset) # this is the default (!)
    return torch.nn.Sequential(normalize_layer, model)

def get_architecture_center_layer(arch: str, dataset: str) -> torch.nn.Module:
    """ Return a neural network (with random weights)
    :param arch: the architecture - should be in the ARCHITECTURES list above
    :param dataset: the dataset - should be in the datasets.DATASETS list
    :return: a Pytorch module
    """
    if arch == "resnet50" and dataset == "imagenet":
        model = torch.nn.DataParallel(resnet50(pretrained=False)).cuda()
        cudnn.benchmark = True
    elif arch == "cifar_resnet20":
        model = resnet_cifar(depth=20, num_classes=10).cuda()
    elif arch == "cifar_resnet110":
        model = resnet_cifar(depth=110, num_classes=10).cuda()
    normalize_layer = get_input_center_layer(dataset) # added for certain SmoothAdv models
    return torch.nn.Sequential(normalize_layer, model)