Upload convnets.py

convnets.py

@@ -0,0 +1,79 @@
+"""
+CNN models for binary and multi-class classifications
+"""
+import torch
+from torch import nn
+
+
+class Convnet(nn.Module):
+    """
+    Convolutional Neural Network for binary classification
+
+    input args: n_classes (int) --> number of classes
+
+    Input shape: [1, 60, 60]
+    
+    Matrix shape (Conv layer):
+
+    Input shape: [N, C_in, H, W]
+        - N: batch_size
+        - C_in: number of input channels
+        - H: height of input planes
+        - W: width of input planes
+
+    - Conv2d(1, 64, (5, 3), 1) --> [64, 56, 58]
+    - MaxPool2d(kernel_size=(2, 1)) --> [64, 28, 58]
+    - Conv2d(64, 128, (5, 3), 1) --> [128, 24, 56]
+    - MaxPool2d(kernel_size=(2, 1)) --> [128, 12, 56]
+    - Conv2d(128, 256, (5, 3), 1) --> [256, 8, 54]
+    - MaxPool2d(kernel_size=(2, 1)) --> [256, 4, 54]
+    
+    Matrix shape (Fully connected layer):
+    - Linear(256 * 4 * 54, 1024) --> [1024]
+    - Linear(1024, 512) --> [512]
+    - Linear(512, 128) --> [128]
+    - Linear(128, 64) --> [64]
+    - Linear(64, n_classes) --> [n_classes]
+
+    Softmax() --> to probability
+    """
+    def __init__(self, n_classes: int) -> None:
+        super().__init__()
+        self.cnn = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(5, 3), stride=1),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(negative_slope=0.01),
+            nn.MaxPool2d(kernel_size=(2, 1)),
+            nn.Conv2d(64, 128, (5, 3), 1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(negative_slope=0.01),
+            nn.MaxPool2d(kernel_size=(2, 1)),
+            nn.Conv2d(128, 256, (5, 3), 1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(negative_slope=0.01),
+            nn.MaxPool2d(kernel_size=(2, 1)),
+        )
+        self.dropout = nn.Sequential(nn.Dropout(0.5))
+        self.fc = nn.Sequential(
+            nn.Linear(256 * 4 * 54, 1024),
+            nn.Linear(1024, 512),
+            nn.Linear(512, 128),
+            nn.Linear(128, 64),
+            nn.Linear(64, n_classes),
+            nn.Softmax()
+        )
+        for layer in self.cnn:
+            if isinstance(layer, nn.Conv2d):
+                nn.init.xavier_normal_(layer.weight)
+                nn.init.constant_(layer.bias, 0.0)
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        forward prop
+        """
+        x = self.cnn(x)
+        x = self.dropout(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x