class DeepCNN(nn.Module):
def __init__(self, num_classes=5):
super(DeepCNN, self).__init__()
self.cnn_layers = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1), # Conv Layer 1
nn.ReLU(),
nn.BatchNorm2d(32),
nn.MaxPool2d(kernel_size=2, stride=2), # Downsampling
nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), # Conv Layer 2
nn.ReLU(),
nn.BatchNorm2d(64),
nn.MaxPool2d(kernel_size=2, stride=2), # Downsampling
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), # Conv Layer 3 (new)
nn.ReLU(),
nn.BatchNorm2d(128),
nn.MaxPool2d(kernel_size=2, stride=2), # Downsampling
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), # Conv Layer 4 (new)
nn.ReLU(),
nn.BatchNorm2d(256),
nn.MaxPool2d(kernel_size=2, stride=2), # Downsampling
)
self.fc_layers = nn.Sequential(
nn.Flatten(),
nn.Linear(256 * 8 * 8, 512), # Fully connected layer
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(512, num_classes)
)
def forward(self, x):
x = self.cnn_layers(x)
x = self.fc_layers(x)
return x
Comments