Changed model

app.py

@@ -23,7 +23,7 @@ model  = create_custom_model()
 # Load saved weights
 model.load_state_dict(
     torch.load(
-        f="./cnn-custom-model-version-3.pt",
+        f="./cnn-custom-model-version-5.pt",
         map_location=torch.device("cpu"),  # load to CPU
     )
 )

cnn-custom-model-version-4.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:520294eec0d142912af0081c04a2c96046cc0f3dae5c986813a018593be15264
+size 43115

cnn-custom-model-version-5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6ed267c84de62f7ee86a28bf3816fb5bf70327314d3197f780e4d44a1bf3355
+size 43115

model.py

@@ -5,6 +5,44 @@ import torchvision.models as models
 from torch import nn
 
 
+# class ConvModel(nn.Module):
+#     def __init__(self, input_shape, output_shape, hidden_units):
+#         super().__init__()
+        
+#         self.block_1 = nn.Sequential(
+#                                 nn.Conv2d(in_channels = input_shape,
+#                                          out_channels = hidden_units,
+#                                          kernel_size = 3,
+#                                          stride = 1,
+#                                          padding = 1),
+#                                 nn.ReLU(),
+#                                 nn.Conv2d(in_channels = hidden_units,
+#                                         out_channels = hidden_units,
+#                                         kernel_size = 3,
+#                                         stride = 1,
+#                                         padding = 1),
+#                                 nn.ReLU(),
+#                                 nn.MaxPool2d(kernel_size = 2,
+#                                             stride = 2))
+        
+#         self.block_2 = nn.Sequential(
+#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
+#                                 nn.ReLU(),
+#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
+#                                 nn.ReLU(),
+#                                 nn.MaxPool2d(2))
+        
+#         self.classifier = nn.Sequential(
+#                                 nn.Flatten(),
+#                                 nn.Linear(in_features = 40960, out_features = output_shape))
+        
+        
+#     def forward(self, x):
+#         x = self.block_1(x)
+#         x = self.block_2(x)
+#         x = self.classifier(x)
+#         return x
+    
 class ConvModel(nn.Module):
     def __init__(self, input_shape, output_shape, hidden_units):
         super().__init__()
@@ -23,27 +61,51 @@ class ConvModel(nn.Module):
                                         padding = 1),
                                 nn.ReLU(),
                                 nn.MaxPool2d(kernel_size = 2,
-                                            stride = 2))
+                                            stride = 2),
+#                                 nn.BatchNorm2d()
+        )
         
         self.block_2 = nn.Sequential(
                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
                                 nn.ReLU(),
                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
                                 nn.ReLU(),
-                                nn.MaxPool2d(2))
+                                nn.MaxPool2d(2),
+#                                 nn.BatchNorm2d()
+        )
+        
+        self.block_3 = nn.Sequential(
+                                nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
+                                nn.ReLU(),
+                                nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
+                                nn.ReLU(),
+                                nn.MaxPool2d(2),
+#                                 nn.BatchNorm2d()
+        )
+        
+        self.block_4 = nn.Sequential(
+                                nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
+                                nn.ReLU(),
+                                nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
+                                nn.ReLU(),
+                                nn.MaxPool2d(2),
+#                                 nn.BatchNorm2d()
+        )
         
         self.classifier = nn.Sequential(
                                 nn.Flatten(),
-                                nn.Linear(in_features = 40960, out_features = output_shape))
+                                nn.Linear(in_features = 5120, out_features = output_shape))
         
         
     def forward(self, x):
         x = self.block_1(x)
         x = self.block_2(x)
+        x = self.block_3(x)
+        x = self.block_4(x)
+#         print(x.shape)
         x = self.classifier(x)
         return x
     
-
 def create_custom_model(seed:int=42):
     """Creates an EfficientNetB2 feature extractor model and transforms.