update model arch

seizure_detection.py

@@ -3,37 +3,31 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
-class SeizureDetector(nn.Module):
-    def __init__(self, num_classes=2):
-        super(SeizureDetector, self).__init__()
-        self.conv1= nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) # 32, 224, 224
+class SeizureDetectionCNN(nn.Module):
+    def init(self, num_classes=2):
+        super(SeizureDetectionCNN, self).init()
+        self.conv1= nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) # 32, 32, 32
 
-        self.pool= nn.MaxPool2d(kernel_size=2, stride=2) # 32, 112, 112
+        self.pool= nn.MaxPool2d(kernel_size=2, stride=2) # 32, 16, 16
 
-        self.conv2= nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) # 64, 112, 112 -> 64, 56, 56
-        self.conv3= nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1) # 128, 56, 56 -> 128, 28, 28
-        self.conv4= nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1) # 256, 28, 28 -> 256, 14, 14
+        self.conv2= nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) # 64, 16, 16 -> 64, 8, 8
 
         # Adding Batch Normalization
         self.bn1 = nn.BatchNorm2d(32)
         self.bn2 = nn.BatchNorm2d(64)
-        self.bn3 = nn.BatchNorm2d(128)
-        self.bn4 = nn.BatchNorm2d(256)
 
         self.dropout = nn.Dropout(p=0.5)  # Dropout with a probability of 50%
 
-        self.fc1= nn.Linear(256*14*14, 120)
+        self.fc1= nn.Linear(64 * 8 * 8, 120)
         self.fc2= nn.Linear(120, 32)
         self.fc3= nn.Linear(32, num_classes)
 
     def forward(self, x):
-        x = self.pool(F.relu(self.bn1(self.conv1(x))))  # 32, 112, 112
-        x = self.pool(F.relu(self.bn2(self.conv2(x))))  # 64, 56, 56
-        x = self.pool(F.relu(self.bn3(self.conv3(x))))  # 128, 28, 28
-        x = self.pool(F.relu(self.bn4(self.conv4(x))))  # 256, 14, 14
+        x = self.pool(F.relu(self.bn1(self.conv1(x))))  # 32, 32, 32
+        x = self.pool(F.relu(self.bn2(self.conv2(x))))  # 64, 8, 8
 
         x = torch.flatten(x, 1)
         x = self.dropout(F.relu(self.fc1(x)))  # Apply dropout
         x = self.dropout(F.relu(self.fc2(x)))  # Apply dropout
         x = self.fc3(x)
-        return x
+        return x