Upload model

modeling_spice_cnn.py

@@ -11,36 +11,25 @@ class SpiceCNNModelForImageClassification(PreTrainedModel):
     def __init__(self, config: SpiceCNNConfig):
         super().__init__(config)
         layers = [
-            nn.Conv2d(
-                config.in_channels,
-                32,
-                kernel_size=3,
-            ),
-            nn.ReLU(),
-            nn.BatchNorm2d(32),
-            nn.Conv2d(32, 32, kernel_size=3),
+            nn.Conv2d(config.in_channels, 16, kernel_size=config.kernel_size, padding=1),
+            nn.BatchNorm2d(16),
             nn.ReLU(),
+            nn.MaxPool2d(kernel_size=config.pooling_size),
+            
+            nn.Conv2d(16, 32, kernel_size=config.kernel_size, padding=1),
             nn.BatchNorm2d(32),
-            nn.Conv2d(32, 32, kernel_size=5, stride=2),
-            nn.ReLU(),
-            nn.Dropout(0.4),
-
-            nn.Conv2d(32, 64, kernel_size=3),
-            nn.ReLU(),
-            nn.BatchNorm2d(64),
-            nn.Conv2d(64, 64, kernel_size=3),
             nn.ReLU(),
+            nn.MaxPool2d(kernel_size=config.pooling_size),
+            
+            nn.Conv2d(32, 64, kernel_size=config.kernel_size, padding=1),
             nn.BatchNorm2d(64),
-            nn.Conv2d(64, 64, kernel_size=5, stride=2),
             nn.ReLU(),
-            nn.BatchNorm2d(64),
-            nn.Dropout(0.4),
-            
-            nn.Flatten(),
-            nn.BatchNorm1d(64),
+            nn.MaxPool2d(kernel_size=2),
+
+            nn.Linear(64 * 4 * 4, 128),
             nn.ReLU(),
-            nn.Dropout(0.4),
-            nn.Linear(64, config.num_classes),
+            nn.Dropout(0.5),
+            nn.Linear(128, config.num_classes)
         ]
         self.model = nn.Sequential(*layers)
 

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:46f937126b04976c9954c667a252be6c1ae0e7283a81c620e8ed53d50d7dcded
-size 792757
+oid sha256:93eef51c610e5fe5b2e79ceb111c0a2b090f38167f859a2394a1cd65af2bef78
+size 632188