Spaces:
Paused
Paused
File size: 5,829 Bytes
39bf8ff |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
from torchvision import models
from torch import nn
import torch
from collections import OrderedDict
def set_parameter_requires_grad(model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False
def initialize_existing_models(model_name, model_type, num_classes, feature_extract, hidden_units, use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet18":
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Sequential(
nn.Linear(num_ftrs, num_classes),
nn.LogSoftmax(dim=1))
input_size = 224
elif model_name == "alexnet":
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Sequential(
nn.Linear(num_ftrs,num_classes),
nn.LogSoftmax(dim=1))
input_size = 224
elif model_name in ["vgg11_bn", "vgg13", "vgg16"]:
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Sequential(
nn.Linear(num_ftrs,num_classes),
nn.LogSoftmax(dim=1))
input_size = 224
elif model_name == "squeezenet":
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Sequential(
nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1)),
nn.LogSoftmax(dim=1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet121":
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Sequential(
nn.Linear(num_ftrs, num_classes),
nn.LogSoftmax(dim=1))
input_size = 224
elif model_name == "inception": # This model expects (299,299) sized images and has auxiliary output
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Sequential(
nn.Linear(num_ftrs, num_classes),
nn.LogSoftmax(dim=1))
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Sequential(
nn.Linear(num_ftrs,num_classes),
nn.LogSoftmax(dim=1))
input_size = 299
else:
print("Invalid model name, please use one of the models supported by this application, exiting...")
exit()
return model_ft, input_size
#Get pre-trained model specifications and override with classifier portion with user activation units
def build_custom_models(model_name, model_type, num_classes, feature_extract, hidden_units, use_pretrained=True):
model_ft = getattr(models, model_name)(pretrained = use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
if model_name == 'resnet18':
in_features = model_ft.fc.in_features
else:
try: #Is there an iterable classifier layer for the model chosen?
iter(model_ft.classifier)
except TypeError: #If no, choose the classifier layer with no index
in_features = model_ft.classifier.in_features
else:
try: #If yes, check if first index has in_features attribute
in_features = model_ft.classifier[0].in_features
except AttributeError: #If No, check if second index has in_features attribute
in_features = model_ft.classifier[1].in_features
hidden_layers = [in_features] + hidden_units
layer_builder = (
lambda i, v : (f"fc{i}", nn.Linear(hidden_layers[i-1], v)),
lambda i, v: (f"relu{i}", nn.ReLU()),
lambda i, v: (f"drop{i}", nn.Dropout())
)
layers = [f(i, v) for i, v in enumerate(hidden_layers) if i > 0 for f in layer_builder]
layers += [('fc_final', nn.Linear(hidden_layers[-1], num_classes)),
('output', nn.LogSoftmax(dim=1))]
if model_name == 'resnet18':
fc = nn.Sequential(OrderedDict(layers))
model_ft.fc = fc
else:
classifier = nn.Sequential(OrderedDict(layers))
model_ft.classifier = classifier
# print("AFTER")
# print(model.classifier)
return model_ft
#Define model/ neural network class
# class ImageClassifier(nn.Module):
# def __init__(self):
# super(ImageClassifer, self).__init__()
# self.flatten = nn.Flatten()
# self.model_stack = nn.Sequential(
# nn.Linear(),
# nn.ReLU(),
# nn.Dropout(0.2),
# nn.Linear(),
# nn.LogSoftmax(dim=1)
# )
# def forward(self, x):
# x = self.flatten(x)
# logits = self.model_stack(x)
# return logits
|