File size: 2,447 Bytes
505b719 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
"""
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
|