def __init__(self):
super(Dunet, self).__init__()
vgg13 = models.vgg13(pretrained=True)
self.conv1 = vgg13.features[0]
self.conv2 = vgg13.features[2]
self.conv3 = vgg13.features[5]
self.conv4 = vgg13.features[7]
self.conv5 = vgg13.features[10]
self.conv6 = vgg13.features[12]
self.dilate_center = Dblock(512)
self.up3 = self.conv_stage(512, 256)
self.up2 = self.conv_stage(256, 128)
self.up1 = self.conv_stage(128, 64)
self.trans3 = self.upsample(512, 256)
self.trans2 = self.upsample(256, 128)
self.trans1 = self.upsample(128, 64)
self.conv_last = nn.Sequential(
nn.Conv2d(64, 1, 3, 1, 1),
nn.Sigmoid()
)
self.max_pool = nn.MaxPool2d(2)
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
if m.bias is not None:
m.bias.data.zero_()
def vgg13(num_classes=1000, pretrained='imagenet'):
"""VGG 13-layer model (configuration "B")
"""
model = models.vgg13(pretrained=False)
if pretrained is not None:
settings = pretrained_settings['vgg13'][pretrained]
model = load_pretrained(model, num_classes, settings)
return model
def getNetwork(args):
if (args.net_type == 'alexnet'):
net = models.alexnet(pretrained=args.finetune)
file_name = 'alexnet'
elif (args.net_type == 'vggnet'):
if (args.depth == 11):
net = models.vgg11(pretrained=args.finetune)
elif (args.depth == 13):
net = models.vgg13(pretrained=args.finetune)
elif (args.depth == 16):
net = models.vgg16(pretrained=args.finetune)
elif (args.depth == 19):
net = models.vgg19(pretrained=args.finetune)
else:
print(
'Error : VGGnet should have depth of either [11, 13, 16, 19]')
sys.exit(1)
file_name = 'vgg-%s' % (args.depth)
elif (args.net_type == 'densenet'):
if (args.depth == 121):
net = models.densenet121(pretrained=args.finetune)
elif (args.depth == 161):
net = models.densenet161(pretrained=args.finetune)
elif (args.depth == 169):
net = models.densenet169(pretrained=args.finetune)
file_name = 'densenet-%s' % (args.depth)
elif (args.net_type == 'resnet'):
net = networks.resnet(args.finetune, args.depth)
file_name = 'resnet-%s' % (args.depth)
elif (args.net_type == 'xception'):
net = pretrainedmodels.xception(num_classes=1000,
pretrained='imagenet')
file_name = 'xception'
elif (args.net_type == 'inception'):
net = models.inception_v3(num_classes=1000, pretrained=args.finetune)
file_name = 'inception'
else:
print(
'Error : Network should be either [alexnet / vggnet / resnet / densenet]'
)
sys.exit(1)
return net, file_name
def build_network(fc_model, trained_model, dropout=0.5, device='cuda'):
model = None
if (trained_model == 'vgg16'):
model = models.vgg16(pretrained=True)
elif (trained_model == 'resnet18'):
model = models.resnet18(pretrained=True)
elif (trained_model == 'vgg13'):
model = models.vgg13(pretrained=True)
for param in model.parameters():
param.requires_grad = False
classifier = nn.Sequential(fc_model)
model.classifier = classifier
model.to(device)
print("Done - {} Network with Classifier... ".format(trained_model))
return model
def create_base_model(arch, hidden_units, class_to_idx):
print(arch)
if arch =='vgg13':
model = models.vgg13(pretrained=True)
elif arch == 'vgg16':
model = models.vgg16(pretrained=True)
elif arch == 'vgg19':
model = models.vgg19(pretrained=True)
else:
raise AssertionError('Architecture not recognized: choose among vgg13, vgg16 and vgg19')
model.classifier = nn.Sequential(OrderedDict([
('fc1', nn.Linear(25088, hidden_units)),
('relu', nn.ReLU()),
('fc2', nn.Linear(hidden_units, 102)),
('output', nn.LogSoftmax(dim=1))
]))
model.class_to_idx = class_to_idx
return model
def checkpoint_loading(filename):
check = torch.load(filename)
model = models.vgg13(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.class_to_idx = check['class_to_idx']
classifier = nn.Sequential(OrderedDict([ ('fc1', nn.Linear(25088, 2960, bias=True)),
('Relu1', nn.ReLU()),
('Dropout1', nn.Dropout(p = 0.5)),
('fc2', nn.Linear(2960, 102, bias=True)),
('output', nn.LogSoftmax(dim=1))
]))
model.classifier = classifier
model.load_state_dict(checkpoint['state_dict'])
return model
def load_and_rebuild_model(nn_checkpoint):
checkpoint = torch.load(nn_checkpoint)
arch = checkpoint['arch']
if arch == 'vgg13':
model = models.vgg13(pretrained=True)
elif arch == 'alexnet':
model = models.alexnet(pretrained=True)
else:
print(
"{} is not a correct architecture. Please choose either vgg13 or alexnet."
.format(arch))
model.classifier = checkpoint['classifier']
model.load_state_dict = checkpoint['state_dict']
model.class_to_idx = checkpoint['class_to_idx']
return model
def load_checkpoint(filepath):
checkpoint = torch.load(filepath)
#Checkpoint
if checkpoint['arch'] == 'vgg13':
model = models.vgg13(pretrained=True)
else:
model = models.vgg16(pretrained=True)
#Freezing our fetures grediance
for param in model.parameters():
param.requires_grad = False
#from checkpoint
model.classifier = checkpoint['classifier']
model.load_state_dict(checkpoint['state_dict'])
model.class_to_idx = checkpoint['class_to_idx']
return model
def load_checkpoint(filepath):
checkpoint = torch.load(filepath)
arch = checkpoint['arch']
model_vgg13 = models.vgg13(pretrained=True)
model_densenet121 = models.densenet121(pretrained=True)
aval_models = {'vgg13': model_vgg13, 'densenet121': model_densenet121}
model = aval_models[arch]
for param in model.parameters():
param.requires_grad=False
model.classifier = checkpoint['classifier'] # NOTE: we need to checkpoint our classfier as well, otherwise an error of mismatch classifier will occur
model.load_state_dict(checkpoint['state_dict'])
model.class_to_idx = checkpoint['class_to_idx']
return model
def model_arch(arch, hidden_units):
if arch == 'vgg16':
model = models.vgg16(pretrained=True)
else:
model = models.vgg13(pretrained=True)
for parameters in model.parameters():
parameters.requires_grad = False
classifier = nn.Sequential(nn.Linear(25088, hidden_units),
nn.ReLU(inplace=True), nn.Dropout(p=0.2),
nn.Linear(hidden_units, hidden_units),
nn.ReLU(inplace=True), nn.Dropout(p=0.2),
nn.Linear(hidden_units, 102),
nn.LogSoftmax(dim=1))
model.classifier = classifier
return model
def InitializeModel(architecture):
switcher = {
'vgg13': models.vgg13(pretrained=True),
'vgg16': models.vgg16(pretrained=True),
'vgg19': models.vgg19(pretrained=True),
'resnet18': models.resnet18(pretrained=True),
'resnet34': models.resnet34(pretrained=True),
'resnet50': models.resnet50(pretrained=True),
'resnet101': models.resnet101(pretrained=True),
'resnet152': models.resnet152(pretrained=True),
'squeezenet1.0': models.squeezenet1_0(pretrained=True),
'squeezenet1.1': models.squeezenet1_1(pretrained=True),
'densenet121': models.densenet121(pretrained=True),
'densenet169': models.densenet169(pretrained=True),
'densenet161': models.densenet161(pretrained=True),
'densenet201': models.densenet201(pretrained=True)
}
model = switcher.get(architecture, "Invalid argument for architecture")
return model
def load_model(architecture, device, hidden_units):
if architecture == 'vgg16':
model = models.vgg16(pretrained=True)
elif architecture == 'vgg13':
model = models.vgg13(pretrained=True)
for param in model.parameters():
param.requires_grad = False
classifier = nn.Sequential(
OrderedDict([('fc1', nn.Linear(25088, 12544)), ('relu1', nn.ReLU()),
('fc2', nn.Linear(12544, hidden_units)),
('relu2', nn.ReLU()),
('fc3', nn.Linear(hidden_units, 102)),
('output', nn.LogSoftmax(dim=1))]))
model.classifier = classifier
model.to(device)
return model
def nn_network(structure='vgg13',dropout=0.5, hidden_layer1 = 120,lr = 0.001,power='gpu'):
'''
Arguments: The architecture for the network(alexnet,densenet121,vgg13), the hyperparameters for the network (hidden layer 1 nodes, dropout and learning rate) and whether to use gpu or not
Returns: The set up model, along with the criterion and the optimizer fo the Training
'''
if structure == 'vgg13':
model = models.vgg13(pretrained=True)
elif structure == 'vgg16':
model = models.vgg16(pretrained=True)
elif structure == 'densenet121':
model = models.densenet121(pretrained=True)
elif structure == 'alexnet':
model = models.alexnet(pretrained = True)
else:
print("Im sorry but {} is not a valid model.Did you mean vgg13,vgg16,densenet121,or alexnet?".format(structure))
for param in model.parameters():
param.requires_grad = False
from collections import OrderedDict
classifier = nn.Sequential(OrderedDict([
('dropout',nn.Dropout(dropout)),
('inputs', nn.Linear(arch[structure], hidden_layer1)),
('relu1', nn.ReLU()),
('hidden_layer1', nn.Linear(hidden_layer1, 90)),
('relu2',nn.ReLU()),
('hidden_layer2',nn.Linear(90,80)),
('relu3',nn.ReLU()),
('hidden_layer3',nn.Linear(80,102)),
('output', nn.LogSoftmax(dim=1))
]))
model.classifier = classifier
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr )
if torch.cuda.is_available() and power == 'gpu':
model.cuda()
return model, criterion, optimizer
def get_model(model_name):
if model_name == 'alexnet':
model = models.alexnet(pretrained=True)
elif model_name == 'vgg19_bn':
model = models.vgg19_bn(pretrained=True)
elif model_name == 'vgg16':
model = models.vgg16(pretrained=True)
elif model_name == 'vgg13':
model = models.vgg13(pretrained=True)
elif model_name == 'vgg11':
model = models.vgg11(pretrained=True)
else:
print("\nError: selecting default vgg16\n")
model = models.vgg16(pretrained=True)
for param in model.parameters():
param.requires_grad = False
return model
def create_model(hidden_units, input_size, output_size, dropout_rate,
selected_model):
print('Creating Model...')
#Load in VGG Model by default
if (selected_model == 'vgg13'):
model = models.vgg13(pretrained=True) #input 25088
input_size = 25088
if (selected_model == 'alexnet'):
model = models.alexnet(pretrained=True) #input 9216
input_size = 9216
#Freeze parameters
for param in model.parameters():
param.requires_grad = False
#creates the OrderedDict to be used in the classifer creation based on initial parameters
classifier_struct = OrderedDict([('fc0',
nn.Linear(input_size, hidden_units[0])),
('dropout0', nn.Dropout(dropout_rate)),
('relu0', nn.ReLU())])
for i in range(len(hidden_units)):
if i == (len(hidden_units) - 1):
classifier_struct[('fc' + str(i + 1))] = nn.Linear(
int(hidden_units[i]), output_size)
classifier_struct['output'] = nn.LogSoftmax(dim=1)
else:
classifier_struct[('fc' + str(i + 1))] = nn.Linear(
int(hidden_units[i]), hidden_units[i + 1])
classifier_struct[('dropout' +
str(i + 1))] = nn.Dropout(dropout_rate)
classifier_struct[('relu' + str(i + 1))] = nn.ReLU()
#Create replacment classifer
classifier = nn.Sequential(classifier_struct)
#Create the model & load in the new classifer
model.classifier = classifier
print('Creating Model...Finished')
return model
def creat_model(pre_model_arch, hidden_layers_size, input_size, output_size):
# Create the network, adjust classifier define the criterion and optimizer
if pre_model_arch == "vgg13":
model_flower_imgs = models.vgg13(pretrained=True)
elif pre_model_arch == "vgg16":
model_flower_imgs = models.vgg16(pretrained=True)
else:
assert 0, "Architecture isn't supported"
# Freeze parameters so we don't backprop through them
for param in model_flower_imgs.parameters():
param.requires_grad = False
# Update the classifier by desired conf.
model_flower_imgs.classifier = creat_classifier(input_size, output_size,
hidden_layers_size)
return model_flower_imgs
def build_model(arch, inputs, hidden_units, output, rate):
"""
This function builds the model to be used in training or predition using
the parameters passed
Params:
arch - Model architecture/type to be used
hidden_units - Number of hidden units to be used in classifier
Returns:
model - Model with replace classifier and freezed gradients
"""
fc1_input = inputs if inputs else 25088
fc2_output = output if output else 102
dp_rate = rate if rate else 0.2
#Check if arch is valid and load the correct model
model_options = ["vgg11", "vgg13", "vgg19"]
if arch in model_options:
print("Building model ---> arch: {}".format(arch))
if arch == "vgg19":
model = models.vgg19(pretrained=True)
elif arch == "vgg13":
model = models.vgg13(pretrained=True)
else:
model = models.vgg11(pretrained=True)
else:
print("Invalid model: {} --> Using arch: vgg11".format(arch))
model = models.vgg11(pretrained=True)
#Freezes gradients
for param in model.parameters():
param.requires_grad = False
#Build classifier and replace it in model
classifier = nn.Sequential(
OrderedDict([("fc1", nn.Linear(fc1_input, hidden_units)),
("relu", nn.ReLU()), ("drop", nn.Dropout(p=dp_rate)),
("fc2", nn.Linear(hidden_units, fc2_output)),
("output", nn.LogSoftmax(dim=1))]))
model.classifier = classifier
return model
def build_network(hidden_layers, drop_p, model='vgg19'):
if model == 'vgg11':
model = models.vgg11(pretrained=True)
elif model == 'vgg13':
model = models.vgg13(pretrained=True)
elif model == 'vgg16':
model = models.vgg16(pretrained=True)
elif model == 'vgg19':
model = models.vgg19(pretrained=True)
else:
print("Please load a VGG network")
return False
# Freeze parameters so we don't backprop through them
for param in model.parameters():
param.requires_grad = False
# Add the classifier
# Started from the lesson, but I tried to adapt to a large variety of hidden layers
classifier_data = OrderedDict()
for i in range(len(hidden_layers)):
classifier_data['dropout' + str(i + 1)] = nn.Dropout(drop_p)
if (i == 0):
classifier_data['fc' + str(i + 1)] = nn.Linear(
25088, hidden_layers[i])
elif (i < len(hidden_layers)):
classifier_data['fc' + str(i + 1)] = nn.Linear(
hidden_layers[i - 1], hidden_layers[i])
else:
classifier_data['fc' + str(i + 1)] = nn.Linear(
hidden_layers[i], 102)
#classifier_data['fc' + str(i+1)] = nn.Linear(hidden_layers[i], 15)
if (i < len(hidden_layers)):
classifier_data['relu' + str(i + 1)] = nn.ReLU()
# add softmax layer
classifier_data['output'] = nn.LogSoftmax(dim=1)
classifier = nn.Sequential(classifier_data)
model.classifier = classifier
return model
def create_model(arch, hidden_features):
if arch == "vgg11":
model = models.vgg11(pretrained=True)
elif arch == "vgg11_bn":
model = models.vgg11_bn(pretrained=True)
elif arch == "vgg13":
model = models.vgg13(pretrained=True)
elif arch == "vgg13_bn":
model = models.vgg13_bn(pretrained=True)
elif arch == "vgg16":
model = models.vgg16(pretrained=True)
elif arch == "vgg16_bn":
model = models.vgg16_bn(pretrained=True)
elif arch == "densenet121":
model = models.densenet121(pretrained=True)
elif arch == "densenet161":
model = models.densenet161(pretrained=True)
else:
print("architecture not supported.")
exit()
# turn off gradient back propagation
for param in model.parameters():
param.requires_grad = False
if arch in ["vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn"]:
in_features = model.classifier[0].in_features
elif arch in ["densenet121", "densenet161"]:
in_features = model.classifier.in_features
if hidden_features < 102 or hidden_features > in_features:
print("please choose a number between {in_features} and 102")
exit()
classifier = nn.Sequential(nn.Linear(in_features, hidden_features),
nn.ReLU(), nn.Dropout(p=0.2),
nn.Linear(hidden_features, 102),
nn.LogSoftmax(dim=1))
model.classifier = classifier
return (model)
def load_model_from_checkpoint(checkpoint_filepath):
checkpoint = torch.load(checkpoint_filepath)
if checkpoint['model_arch'] == 'vgg19':
model = models.vgg19(pretrained=True)
arch = 'vgg19'
elif checkpoint['model_arch'] == 'vgg13':
model = models.vgg13(pretrained=True)
arch = 'vgg13'
elif checkpoint['model_arch'] == 'vgg16':
model = models.vgg16(pretrained=True)
arch = 'vgg16'
else:
print('The models architecture is not recognized')
for param in model.parameters():
param.requires_grad = False
model.class_to_idx = checkpoint['class_to_idx']
model.classifier = Classifier(checkpoint['hidden_units'])
model.load_state_dict(checkpoint['state_dict'])
#model.to(checkpoint['model_device'])
return model
def Network(hidden_layers):
# Check architecture
if args.arch == 'vgg16':
model = models.vgg16(pretrained=True)
else:
model = models.vgg13(pretrained=True)
# Freeze parameters
for param in model.parameters():
param.requires_grad = False
# Defining classifier
classifier = nn.Sequential(
OrderedDict([('fc1', nn.Linear(25088, hidden_layers[0])),
('relu', nn.ReLU()), ('drop', nn.Dropout(p=0.5)),
('fc2', nn.Linear(hidden_layers[0], hidden_layers[1])),
('relu2', nn.ReLU()), ('drop2', nn.Dropout(p=0.5)),
('fc3', nn.Linear(hidden_layers[1], 102)),
('output', nn.LogSoftmax(dim=1))]))
# Replacing classifier
model.classifier = classifier
return model
def load_checkpoint(checkpoint_path):
checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc:storage)
model = models.vgg13(pretrained=True)
for param in model.parameters():
param.requires_grad=False
for param in model.classifier.parameters():
param.requires_grad= True
model.class_to_idx = checkpoint['class_to_idx']
model.classifier = nn.Sequential(
nn.Linear(25088, 4096),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(4096, 102),
nn.LogSoftmax(dim=1))
model.load_state_dict(checkpoint['model_state_dict'], strict=False)
model.eval()
return model
def build_network(self):
model_select = {
"vgg13": models.vgg13(pretrained=True),
"vgg11": models.vgg11(pretrained=True)
}
model = model_select[self.arch]
print(model)
for param in model.parameters():
param.requires_grad = False
classifier = nn.Sequential(
nn.Linear(model.classifier[0].in_features, self.hidden_layer),
nn.ReLU(), nn.Dropout(0.05),
nn.Linear(self.hidden_layer, self.output_size),
nn.LogSoftmax(dim=1))
model.classifier = classifier
filename = 'model_base.pkl'
pickle.dump(model, open(filename, 'wb'))
return model, classifier
def generate_model(arch, units):
if (arch == "vgg13"):
model = models.vgg13(pretrained=True)
input_size = model.classifier[0].in_features
elif (arch == "vgg16"):
model = models.vgg16(pretrained=True)
input_size = model.classifier[0].in_features
elif (arch == "densenet121"):
model = models.densenet121(pretrained=True)
input_size = model.classifier.in_features
else:
print("error")
exit()
classifier_dict = OrderedDict([("fc1",nn.Linear(input_size, units)),
("relu1", nn.ReLU()),
("dropout", nn.Dropout(0.2)),
("fc2", nn.Linear(units,102)),
("p_out", nn.LogSoftmax(dim=1))])
model.classifier = nn.Sequential(classifier_dict)
return model, classifier_dict
def load_checkpoint():
checkpoint = torch.load("checkpoint.pth")
structure = checkpoint['architecture']
if structure == 'vgg13':
model = models.vgg13(pretrained=True)
model.name = "vgg13"
else:
a = {}
exec("model = models.{}(pretrained=True)".format(structure), globals(),
a)
model = a["model"]
model.name = structure
model.state_dict(checkpoint['state_dict'])
model.classifier = checkpoint['classifier']
model.class_to_idx = checkpoint['class_to_idx']
return model
def load_checkpoint(path):
checkpoint = torch.load(path)
if checkpoint['arch'] == 'vgg13':
model = models.vgg13(pretrained=True)
else: #vgg11
model = models.vgg11(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.classifier = checkpoint['classifier']
model.class_to_idx = checkpoint['class_to_idx']
model.load_state_dict(checkpoint['state_dict'])
optimizer = checkpoint['optimizer']
optimizer.load_state_dict(checkpoint['optimizer_dict'])
model.idx_to_class = {v: k for k, v in model.class_to_idx.items()}
return model
def load_saved_model(filename):
checkpoint_dict = torch.load(filename)
if checkpoint_dict["model_used"] == "vgg11":
model_check = models.vgg11(pretrained=True)
elif checkpoint_dict["model_used"] == "vgg13":
model_check = models.vgg13(pretrained=True)
elif checkpoint_dict["model_used"] == "vgg16":
model_check = models.vgg16(pretrained=True)
elif checkpoint_dict["model_used"] == "vgg19":
model_check = models.vgg19(pretrained=True)
for param in model_check.parameters():
param.requires_grad = False
optimizer = checkpoint_dict["optimizer"]
model_check.classifier = checkpoint_dict["classifier"]
model_check.load_state_dict(checkpoint_dict["state_dict"])
optimizer.load_state_dict(checkpoint_dict["optimizer_dict"])
model_check.class_to_idx = checkpoint_dict["class_to_idx"]
return model_check
def getModel(output_units,
hidden_units,
arch='vgg16',
class_to_idx=None,
learning_rate=1e-3,
criterion=NLLLoss,
optimizer=Adam):
#Choosing between the models
if arch == 'vgg13':
model = vgg13(pretrained=True)
else:
model = vgg16(pretrained=True)
#Turning off all the pretrained parameters, since we don't want to retrain them
for param in model.parameters():
param.requires_grad = False
#We define our new classifier, with input that match the default vgg16 classifier,
#and output equals to the number of classes
classifier = Sequential(
OrderedDict([
('fc1', Linear(25088, hidden_units)),
('act1', ReLU()),
('Dropout1', Dropout(p=0.5)),
('fc2', Linear(hidden_units, hidden_units)),
('act2', ReLU()),
('Dropout2', Dropout(p=0.5)),
('fc3', Linear(hidden_units, output_units)),
('output', LogSoftmax(dim=1)),
]))
#we deattach the default vgg16 classeifer and plug the one
model.classifier = classifier
model.class_to_idx = class_to_idx
criterion = criterion()
optimizer = optimizer(model.classifier.parameters(), lr=learning_rate)
model.criterion = criterion
model.optimizer = optimizer
model.arch = arch
return model
def load_checkpoint(file_path, use_this_device):
print('- Loading checkpoint from file:',file_path)
print('- Using device:', use_this_device)
checkpoint = torch.load(file_path, map_location = use_this_device)
model_arch = checkpoint['model_arch']
if model_arch == 'vgg19':
model = models.vgg19(pretrained=True)
elif model_arch == 'vgg16':
model = models.vgg16(pretrained=True)
elif model_arch == 'vgg13':
model = models.vgg13(pretrained = True)
else:
print('- Model architecture {} unsupported: using vgg19 instead.'.format(model_arch))
model = models.vgg19(pretrained = True)
# Freeze parameters!
for param in model.parameters():
param.requires_grad = False
# Recreate the classifier
from collections import OrderedDict
classifier = nn.Sequential(OrderedDict([
('fc1', nn.Linear(25088, checkpoint['hidden_units'])), # model has 25088 in_features
('relu', nn.ReLU()),
('fc2', nn.Dropout(p = 0.5)), # this is defined in train.py
('fc2', nn.Linear(checkpoint['hidden_units'], 102)), # model has 102 out_features
('output', nn.LogSoftmax(dim=1))
]))
model.classifier = classifier
# Load model state
model.load_state_dict(checkpoint['model_state_dict'])
# Load class to idx
model.class_to_idx = checkpoint['class_to_idx']
return model
def load_model(name,
num_classes,
input_height=64,
input_width=64,
num_of_channels=3):
if name == 'COAPModNet':
net = COAPModNet(num_classes=num_classes)
elif name == 'COAPNet':
net = COAPNet(num_classes=num_classes)
elif name == 'SimpleNet':
net = SimpleNet(num_classes=num_classes)
elif name == 'AlexNet':
net = models.AlexNet(num_classes=num_classes)
elif name == 'PlanktonNet':
net = PlanktonNet(num_classes=num_classes)
elif name == 'ResNet18':
net = models.resnet18(num_of_channels, num_classes)
elif name == 'ResNet34':
net = models.resnet34(num_of_channels, num_classes)
elif name == 'ResNet50':
net = models.resnet50(num_of_channels, num_classes)
elif name == 'ResNet101':
net = models.resnet101(num_of_channels, num_classes)
elif name == 'ResNet152':
net = models.resnet152(num_of_channels, num_classes)
elif name == 'VGGNet11':
net = models.vgg11(num_classes=num_classes)
elif name == 'VGGNet13':
net = models.vgg13(num_classes=num_classes)
elif name == 'VGGNet16':
net = models.vgg16(num_classes=num_classes)
elif name == 'VGGNet19':
net = models.vgg19(num_classes=num_classes)
elif name == 'ResNext50':
net = models.resnext50_32x4d(num_classes=num_classes)
elif name == 'ResNext101':
net = models.resnext101_32x8d(num_classes=num_classes)
elif name == 'GoogLeNet':
net = models.GoogLeNet(num_classes=num_classes)
return net
def create_model(model_name, hidden_units, learning_rate, device):
if model_name == 'vgg13':
model = models.vgg13(pretrained=True)
elif model_name == 'vgg16':
model = models.vgg16(pretrained=True)
else:
model = models.vgg19(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.classifier = nn.Sequential(nn.Linear(25088, hidden_units), nn.ReLU(),
nn.Dropout(0.25),
nn.Linear(hidden_units, 102),
nn.LogSoftmax(dim=1))
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)
model.to(device)
return model, criterion, optimizer
def load_checkpoint(path):
checkpoint = torch.load(path, map_location="cpu")
if checkpoint['arch'] == 'vgg11':
model = models.vgg11(pretrained=True)
elif checkpoint['arch'] == 'vgg13':
model = models.vgg13(pretrained=True)
elif checkpoint['arch'] == 'vgg16':
model = models.vgg16(pretrained=True)
elif checkpoint['arch'] == 'vgg19':
model = models.vgg19(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.features = checkpoint['features']
model.classifier = checkpoint['classifier']
model.class_to_idx = checkpoint['class_to_idx']
model.load_state_dict(state_dict=checkpoint['state_dict'])
return model
