my_project.model

  1"""
  2Defines the neural network architecture for the Fashion-MNIST classifier.
  3
  4This module contains the `Net` class, a PyTorch LightningModule that
  5encapsulates the model's structure (a simple CNN), the forward pass logic,
  6and the training, validation, and test steps.
  7"""
  8
  9import torch
 10import torch.nn as nn
 11import torch.optim as optim
 12import torch.nn.functional as F
 13import pytorch_lightning as pl
 14
 15
 16class Net(pl.LightningModule):
 17    """
 18    Simple convolutional neural network for Fashion-MNIST classification.
 19
 20    Architecture
 21    ------------
 22    - Conv2d(1 → 16, kernel_size=3)
 23    - ReLU
 24    - MaxPool2d(kernel_size=2)
 25    - Flatten
 26    - Linear(16*13*13 → 32)  # Note: This will be configurable
 27    - ReLU
 28    - Linear(32 → 10)       # Note: This will be configurable
 29
 30    Loss: CrossEntropyLoss
 31
 32    Examples
 33    --------
 34    >>> model = Net()
 35    >>> x = torch.randn(8, 1, 28, 28)
 36    >>> out = model(x)
 37    >>> out.shape
 38    torch.Size([8, 10])
 39    """
 40
 41    def __init__(
 42        self,
 43        num_filters: int = 16,
 44        hidden_size: int = 32,
 45        lr: float = 1e-3,
 46    ):
 47        super().__init__()
 48        self.save_hyperparameters()  # Save hyperparameters
 49
 50        self.num_filters = num_filters
 51        self.hidden_size = hidden_size
 52        self.lr = lr
 53
 54        self.conv = nn.Conv2d(1, self.num_filters, 3)
 55        self.pool = nn.MaxPool2d(2)
 56        self.flat = nn.Flatten()
 57        # Calculate the flattened size after conv and pool
 58        self.fc1 = nn.Linear(self.num_filters * 13 * 13, self.hidden_size)
 59        self.fc2 = nn.Linear(self.hidden_size, 10)
 60        self.loss_fn = nn.CrossEntropyLoss()
 61
 62    def forward(self, x):
 63        """
 64        Forward pass of the network.
 65
 66        Parameters
 67        ----------
 68        x : torch.Tensor
 69            Input tensor of shape (N, 1, 28, 28).
 70
 71        Returns
 72        -------
 73        torch.Tensor
 74            Output logits of shape (N, 10).
 75        """
 76
 77        x = self.pool(F.relu(self.conv(x)))
 78        x = self.flat(x)
 79        x = F.relu(self.fc1(x))
 80        return self.fc2(x)
 81
 82    def training_step(self, batch, batch_idx):
 83        """
 84        Training step for a single batch.
 85
 86        Parameters
 87        ----------
 88        batch : tuple
 89            A batch of data (images, labels).
 90        batch_idx : int
 91            Index of the batch.
 92
 93        Returns
 94        -------
 95        torch.Tensor
 96            Training loss.
 97        """
 98
 99        xb, yb = batch
100        out = self(xb)
101        loss = self.loss_fn(out, yb)
102        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
103        return loss
104
105    def validation_step(self, batch, batch_idx):
106        """
107        Validation step for a single batch.
108
109        Parameters
110        ----------
111        batch : tuple
112            A batch of data (images, labels).
113        batch_idx : int
114            Index of the batch.
115
116        Returns
117        -------
118        None
119        """
120
121        xb, yb = batch
122        out = self(xb)
123        preds = out.argmax(1)
124        acc = (preds == yb).float().mean()
125        self.log("val_acc", acc, prog_bar=True)
126
127    def test_step(self, batch, batch_idx):
128        """
129        Test step for a single batch.
130
131        Parameters
132        ----------
133        batch : tuple
134            A batch of data (images, labels).
135        batch_idx : int
136            Index of the batch.
137
138        Returns
139        -------
140        None
141        """
142
143        xb, yb = batch
144        out = self(xb)
145        preds = out.argmax(1)
146        acc = (preds == yb).float().mean()
147        self.log("test_acc", acc, prog_bar=True, on_epoch=True)
148
149    def configure_optimizers(self):
150        """
151        Define optimizer for training.
152
153        Returns
154        -------
155        torch.optim.Optimizer
156            Adam optimizer with default parameters.
157        """
158
159        return optim.Adam(self.parameters(), lr=self.lr)