Defining a model class in PyTorch - ML Experiment: Train & Evaluate

Experiment - Defining a model class

Problem:You want to create a simple neural network model class in PyTorch to classify images from the MNIST dataset.

Current Metrics:Training accuracy: 85%, Validation accuracy: 82%, Training loss: 0.45, Validation loss: 0.50

Issue:The current model is defined using a sequential container only, which limits flexibility for future changes and custom layers.

Your Task

Define a PyTorch model class that inherits from nn.Module with an __init__ method and a forward method to build the same model architecture.

Use nn.Module inheritance.

Include at least one hidden layer with ReLU activation.

Output layer must have 10 units for 10 classes.

Do not change the model architecture or training code.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleMNISTModel(nn.Module):
    def __init__(self):
        super(SimpleMNISTModel, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)  # input layer to hidden
        self.fc2 = nn.Linear(128, 10)     # hidden to output

    def forward(self, x):
        x = x.view(x.size(0), -1)  # flatten images
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Example usage:
# model = SimpleMNISTModel()
# output = model(torch.randn(64, 1, 28, 28))  # batch of 64 images

Created a class SimpleMNISTModel inheriting from nn.Module.

Defined __init__ method to initialize two linear layers.

Defined forward method to flatten input, apply first layer with ReLU, then output layer.

Replaced sequential container with explicit layer definitions for flexibility.

Results Interpretation

Before: Model was defined using nn.Sequential, limiting customization.
After: Model is a class with explicit layers and forward method, allowing easy changes and better understanding.

Training and validation metrics remain the same, confirming the model architecture is unchanged.

Defining a model as a class in PyTorch gives you full control over the model's structure and data flow, which is essential for building more complex or custom models.

Bonus Experiment

Add a dropout layer after the first hidden layer to reduce overfitting.

💡 Hint

Use nn.Dropout in __init__ and apply it in forward after ReLU activation.