model.py

from turtle import forward
from torch import nn

class Mnist(nn.Module):
    """
    A simple convolutional neural network for  digit classification.
    This model consists of two convolutional layers followed by ReLU activations and max pooling,
    a flattening layer, and a fully connected output layer.
    
    Architecture:
        - Conv2d (in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2)
        - ReLU
        - MaxPool2d (kernel_size=2)
        - Conv2d (in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2)
        - ReLU
        - MaxPool2d (kernel_size=2)
        - Flatten
        - Linear (in_features=32*7*7, out_features=10)
    Args:
        *args: Variable length argument list for nn.Module.
        **kwargs: Arbitrary keyword arguments for nn.Module.
    Input shape:
        Tensor of shape (batch_size, 1, 28, 28)
    Output shape:
        Tensor of shape (batch_size, 10), representing class logits for 10 digits (0-9).
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1,out_channels=16,kernel_size=5,stride=1,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16,out_channels=32,kernel_size=5,stride=1,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        
        self.flatten=nn.Flatten(start_dim=1)

        self.fc = nn.Linear(in_features=32*7*7, out_features=10) 

    def forward(self,x):
        x=self.conv1(x)
        x=self.conv2(x)
        x=self.flatten(x)
        x=self.fc(x)
        return x