-
Notifications
You must be signed in to change notification settings - Fork 1
/
architectures.py
43 lines (40 loc) · 1.88 KB
/
architectures.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import torch
# from torchvision.models.resnet import resnet50
import torch.backends.cudnn as cudnn
from archs.cifar_resnet import resnet as resnet_cifar
from archs.cifar_resnet import lenet300
from archs.cifar_resnet import lenet5, lenet_5_caffe
from archs.cifar_resnet import fcn, wide_resnet
from archs.cifar_resnet import vgg19, resnet32, resnet50
from torch.nn.functional import interpolate
ARCHITECTURES = ["resnet50", "lenet300", "lenet5", "vgg19", "resnet32", "vgg16", "lenet_5_caffe"]
def get_architecture(arch: str, dataset: str, device) -> 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)).to(device)
cudnn.benchmark = True
elif arch == "resnet50" and dataset == "tiny_imagenet":
model = resnet50(num_classes=200).to(device)
elif arch == "resnet32" and dataset == "tiny_imagenet":
model = resnet32(num_classes=200).to(device)
elif arch == "resnet32" and dataset == "cifar100":
model = resnet32(num_classes=100).to(device)
elif arch == "resnet32":
model = resnet32(num_classes=10).to(device)
elif arch == "vgg19" and dataset == "tiny_imagenet":
model = vgg19(num_classes=200).to(device)
elif arch == "lenet300":
model = lenet300(num_classes=10).to(device)
elif arch == "lenet5":
model = lenet5(num_classes=10).to(device)
elif arch == "lenet_5_caffe":
model = lenet_5_caffe().to(device)
elif arch == "vgg19" and dataset == "cifar100":
model = vgg19(num_classes=100).to(device)
elif arch == "vgg19":
model = vgg19(num_classes=10).to(device)
return model