Update model.py

model.py

@@ -5,28 +5,36 @@ import torch.nn.functional as F
 class AudioCNN(nn.Module):
     def __init__(self, num_classes, input_length):
         super(AudioCNN, self).__init__()
-        # Ensure input_length is an integer, not a tuple
+
+        # Ensure input_length is integer
         if isinstance(input_length, (tuple, list)):
             input_length = input_length[0]
 
-        # 1D Convolution for raw audio
+        # Convolutional layers
         self.conv1 = nn.Conv1d(1, 16, kernel_size=3, stride=1, padding=1)
-        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
-        
-        # Dynamically compute flattened size
+        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.conv2 = nn.Conv1d(16, 32, kernel_size=3, stride=1, padding=1)
+        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)
+
+        # Compute output size after conv + pooling dynamically
         with torch.no_grad():
-            dummy_input = torch.zeros((1, 1, input_length))  # batch=1, channels=1
-            x = self.pool(F.relu(self.conv1(dummy_input)))
+            dummy_input = torch.zeros(1, 1, input_length)  # batch=1, channel=1
+            x = self.pool1(F.relu(self.conv1(dummy_input)))
+            x = self.pool2(F.relu(self.conv2(x)))
             flattened_size = x.numel() // x.size(0)
-        
-        # Fully connected layer
-        self.fc1 = nn.Linear(flattened_size, num_classes)
+
+        # Fully connected layers
+        self.fc1 = nn.Linear(flattened_size, 128)
+        self.fc2 = nn.Linear(128, num_classes)
 
     def forward(self, x):
-        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool1(F.relu(self.conv1(x)))
+        x = self.pool2(F.relu(self.conv2(x)))
         x = x.view(x.size(0), -1)
-        x = self.fc1(x)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
         return x
 
 
 
+