Upload model

modeling_spice_cnn.py

@@ -1,9 +1,11 @@
 import torch.nn as nn
+
 # from torchsummary import summary
 
 from transformers import PreTrainedModel
 
-from hf_models.models.spice_cnn.configuration_spice_cnn import SpiceCNNConfig
+from .configuration_spice_cnn import SpiceCNNConfig
+
 
 class SpiceCNNModelForImageClassification(PreTrainedModel):
     config_class = SpiceCNNConfig
@@ -11,26 +13,25 @@ class SpiceCNNModelForImageClassification(PreTrainedModel):
     def __init__(self, config: SpiceCNNConfig):
         super().__init__(config)
         layers = [
-            nn.Conv2d(config.in_channels, 16, kernel_size=config.kernel_size, padding=1),
+            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.ReLU(),
             nn.MaxPool2d(kernel_size=config.pooling_size),
-
             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.Linear(64 * 3 * 3, 128),
             nn.ReLU(),
             nn.Dropout(0.5),
-            nn.Linear(128, config.num_classes)
+            nn.Linear(128, config.num_classes),
         ]
         self.model = nn.Sequential(*layers)
 
@@ -41,7 +42,8 @@ class SpiceCNNModelForImageClassification(PreTrainedModel):
             loss = loss_fnc(logits, labels)
             return {"loss": loss, "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:4c6b92faf374dee8dcdbaa0b65dd62f3fde1178284ffe3a8e254b49f3c73e249
+oid sha256:9099bac65074f55cf277404cf2daffecc6893d84cdda384e6e95eb6d6a257914
 size 402812