Upload neural_network_model.py

neural_network_model.py

@@ -0,0 +1,131 @@
+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