Defining a model class in PyTorch

Introduction

We define a model class to create a reusable and clear structure for our machine learning model. This helps us organize layers and how data flows through them.

When you want to build a neural network with multiple layers.

When you need to customize how data moves through your model.

When you want to reuse the same model structure with different data.

When you want to save and load your model easily.

When you want to experiment with different model designs.

Syntax

PyTorch

import torch.nn as nn

class ModelName(nn.Module):
    def __init__(self):
        super(ModelName, self).__init__()
        # Define layers here

    def forward(self, x):
        # Define data flow here
        return x

The class inherits from nn.Module, which is the base for all models in PyTorch.

The __init__ method sets up layers, and forward defines how input moves through them.

Examples

A simple model with one linear layer that changes input size from 10 to 5.

PyTorch

import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

A model with two layers and a ReLU activation between them.

PyTorch

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

class TwoLayerNN(nn.Module):
    def __init__(self):
        super(TwoLayerNN, self).__init__()
        self.layer1 = nn.Linear(8, 4)
        self.layer2 = nn.Linear(4, 2)

    def forward(self, x):
        x = self.layer1(x)
        x = F.relu(x)
        x = self.layer2(x)
        return x

Sample Model

This code defines a simple model with one linear layer and softmax output. It runs the model on two input samples and prints the probabilities.

PyTorch

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

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(3, 2)

    def forward(self, x):
        x = self.fc1(x)
        return F.softmax(x, dim=1)

# Create model instance
model = SimpleModel()

# Sample input tensor with batch size 2 and 3 features
input_tensor = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

# Get model output
output = model(input_tensor)
print(output)

Important Notes

Always call super().__init__() in your model's __init__ method to initialize the base class.

The forward method is where you describe how input data moves through layers.

Use nn.functional for activation functions like relu or softmax inside forward.

Summary

Defining a model class helps organize layers and data flow clearly.

Inherit from nn.Module and implement __init__ and forward methods.

Use the model by creating an instance and passing input tensors to it.