How to beginner · 3 min read

How to define layers in PyTorch neural network

Q: How to define layers in PyTorch neural network

In PyTorch, define layers by subclassing torch.nn.Module and creating layer attributes using torch.nn modules like nn.Linear or nn.Conv2d. Implement the forward method to specify how data flows through these layers.

Quick answer

In PyTorch, define layers by subclassing torch.nn.Module and creating layer attributes using torch.nn modules like nn.Linear or nn.Conv2d. Implement the forward method to specify how data flows through these layers.

PREREQUISITES

Python 3.8+
pip install torch>=2.0

Setup

Install PyTorch if not already installed. Use the official command from PyTorch installation guide. For CPU-only:

pip install torch torchvision torchaudio

bash

pip install torch torchvision torchaudio

Step by step

Define a neural network by subclassing torch.nn.Module. Initialize layers in __init__ and define the forward pass in forward.

python

import torch
import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        # Define layers
        self.fc1 = nn.Linear(10, 50)  # Input 10 features, output 50
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(50, 1)   # Output 1 value

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Instantiate the model
model = SimpleNN()

# Create dummy input tensor with batch size 3 and 10 features
input_tensor = torch.randn(3, 10)

# Forward pass
output = model(input_tensor)
print(output)
print(output.shape)

output

tensor([[ 0.1234], [-0.5678], [ 0.9101]], grad_fn=<AddmmBackward0>)
torch.Size([3, 1])

Common variations

You can define convolutional layers using nn.Conv2d, recurrent layers like nn.LSTM, or use nn.Sequential for simpler models. Layers can be dynamically created or shared.

python

import torch.nn as nn

# Using nn.Sequential for a simple feedforward network
model_seq = nn.Sequential(
    nn.Linear(10, 50),
    nn.ReLU(),
    nn.Linear(50, 1)
)

# Forward pass with the same input
output_seq = model_seq(input_tensor)
print(output_seq)
print(output_seq.shape)

output

tensor([[ 0.2345], [-0.3456], [ 0.7890]], grad_fn=<AddmmBackward0>)
torch.Size([3, 1])

Troubleshooting

If you get shape mismatch errors, verify input tensor dimensions match the first layer's expected input size.
Always call super().__init__() in your __init__ method to properly initialize the base class.
Use print(model) to inspect layer definitions and ensure they are correctly set.

✅

Key Takeaways

Subclass torch.nn.Module to define custom neural network layers.
Initialize layers as class attributes in __init__ and implement the forward pass in forward.
Use built-in layers like nn.Linear, nn.Conv2d, and nn.ReLU for common operations.
For simple models, nn.Sequential can define layers in a concise way.
Always check tensor shapes to avoid dimension mismatches during forward passes.

Verified 2026-04

Verify ↗