Fixed model

app.py

@@ -7,6 +7,7 @@ from sklearn.preprocessing import LabelEncoder
 import torch
 import torch.nn.functional as F
 from torchvision import transforms
+import torchvision
 import torchvision.models as models
 from torchvision.datasets import ImageFolder
 from torch.utils.data.dataset import Dataset
@@ -16,21 +17,17 @@ from sklearn.model_selection import train_test_split
 from tqdm.notebook import tqdm
 
 
-            
-class DogBreedPretrainedWideResnet(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        
-        self.network = models.wide_resnet50_2(pretrained=True)
-        # Replace last layer
-        num_ftrs = self.network.fc.in_features
-        self.network.fc = torch.nn.Sequential(
-            torch.nn.Linear(num_ftrs, 120),
-            torch.nn.LogSoftmax(dim=1)
-        )
-        
-    def forward(self, xb):
-        return self.network(xb)
+class net50(torch.nn.Module):
+    def __init__(self, base_model, base_out_features, num_classes):
+        super(net50,self).__init__()
+        self.base_model=base_model
+        self.linear1 = torch.nn.Linear(base_out_features, 512)
+        self.output = torch.nn.Linear(512,num_classes)
+    def forward(self,x):
+        x = F.relu(self.base_model(x))
+        x = F.relu(self.linear1(x))
+        x = self.output(x)
+        return x
 
 def get_default_device():
     if torch.cuda.is_available():
@@ -61,7 +58,11 @@ def transform_image(image_bytes):
                                             [0.229, 0.224, 0.225])])
     return my_transforms(image_bytes).unsqueeze(0)
 
-model = DogBreedPretrainedWideResnet()
+res = torchvision.models.resnet50(pretrained=True)
+for param in res.parameters():    ## Freezing layers
+    param.requires_grad=False
+
+model = net50(base_model=res, base_out_features=res.fc.out_features, num_classes=120)
 model.load_state_dict(torch.load(PATH,map_location))
 model.eval()
 breeds=['Chihuahua',
@@ -184,6 +185,6 @@ breeds=['Chihuahua',
  'dingo',
  'dhole',
  'African hunting dog']
-iface = gr.Interface(image_mod, gr.Image(type="pil"), "text", examples=["doggo1.jpg","doggo2.jpg","doggo3.png","doggo4.png"])
+iface = gr.Interface(image_mod, gr.Image(type="pil"), "text", examples=["doggo1.png","doggo2.jpg","doggo3.png","doggo4.png"])
 
 iface.launch()

model/model.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c7b67545692a07c4d4e255343f7fa5ff6a327c56ac4bcca38efc2867c8eb8a38
-size 268684285
+oid sha256:6ac024f745a459387e271e7be9ceba0403f19b5caa78138926350b667caad5ff
+size 104830125