update to CNN architecture

src/models/models.py

@@ -72,17 +72,29 @@ def __init__(self):
         super(ConvLayers, self).__init__()
         # a little strange = # of filters, usually goes from small to large
         # double check on architecture decisions
+        '''
         self.conv1 = nn.Conv2d(1, 10, kernel_size=3, padding=1)
         self.conv2 = nn.Conv2d(10, 10, kernel_size=3, padding=1)
         self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2, padding=1)
         self.conv3 = nn.Conv2d(10, 10, kernel_size=3, padding=1)
         self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2, padding=1)
         self.conv4 = nn.Conv2d(10, 5, kernel_size=3, padding=1)
         self.conv5 = nn.Conv2d(5, 5, kernel_size=3, padding=1)
+        '''
+        self.conv1 = nn.Conv2d(1, 5, kernel_size=7, padding=1)
+        self.conv2 = nn.Conv2d(5, 5, kernel_size=7, padding=1)
+        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(5, 5, kernel_size=5, padding=1)
+        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(5, 10, kernel_size=3, padding=1)
+        self.conv5 = nn.Conv2d(10, 10, kernel_size=3, padding=1)
+
         self.flatten = nn.Flatten()
 
     def forward(self, x):
-        # print('input shape', x.shape)
+        assert x.dim() != 2, \
+            f"should enter here with a dimension of at least 3, " \
+            f"{x.dim()}, {x.shape}"
         if x.dim() == 3:  # Check if the input is of shape (batchsize, 32, 32)
             # Add channel dimension, becomes (batchsize, 1, 32, 32)
             x = x.unsqueeze(1)
@@ -124,7 +136,7 @@ def model_setup_DER(loss_type, DEVICE, n_hidden=64, data_type="0D"):
         model = torch.nn.Sequential(
             conv_layers,
             Model(
-                n_hidden=n_hidden, n_input=405, n_output=4
+                n_hidden=n_hidden, n_input=360, n_output=4  # 405
             ),  # Adjust input size according to the flattened output size
             Layer(),
         )
@@ -171,7 +183,7 @@ def model_setup_DE(loss_type, DEVICE, n_hidden=64, data_type="0D"):
         model = torch.nn.Sequential(
             conv_layers,
             Model(
-                n_hidden=n_hidden, n_input=405, n_output=2
+                n_hidden=n_hidden, n_input=360, n_output=2
             ),  # Adjust input size according to the flattened output size
             Layer(),
         )