Upload model

modeling_spice_cnn.py

@@ -1,9 +1,9 @@
 import torch.nn as nn
+# from torchsummary import summary
 
 from transformers import PreTrainedModel
 
-from .configuration_spice_cnn import SpiceCNNConfig
-
+from hf_models.models.spice_cnn.configuration_spice_cnn import SpiceCNNConfig
 
 class SpiceCNNModelForImageClassification(PreTrainedModel):
     config_class = SpiceCNNConfig
@@ -21,7 +21,13 @@ class SpiceCNNModelForImageClassification(PreTrainedModel):
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=config.pooling_size),
 
-            nn.Linear(7*7*32, 128),
+            nn.Conv2d(32, 64, kernel_size=config.kernel_size, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=config.pooling_size),
+
+            nn.Flatten(),
+            nn.Linear(64*3*3, 128),
             nn.ReLU(),
             nn.Dropout(0.5),
             nn.Linear(128, config.num_classes)
@@ -34,4 +40,8 @@ class SpiceCNNModelForImageClassification(PreTrainedModel):
             loss_fnc = nn.CrossEntropyLoss()
             loss = loss_fnc(logits, labels)
             return {"loss": loss, "logits": logits}
-        return {"logits": logits}
+        return {"logits": logits}
+    
+# config = SpiceCNNConfig(in_channels=1)
+# cnn = SpiceCNNModelForImageClassification(config)
+# summary(cnn, (1,28,28))

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1d40410140194428e35927afeb15389e18b894f35bd109991e281597b3938623
-size 833767
+oid sha256:548d7b0aa1d5b69f63f128dfa0d9d343bd9ca053a29c2dae2114d5b6f2d7b5c6
+size 402812