Training a Neural Network in PyTorch

Introduction

Once you’ve defined your model, the next step is training. In PyTorch, training involves preparing data, defining a loss function and optimizer, and implementing the training loop where forward and backward passes are performed.

Step-by-Step Training Workflow

1. Create a dataset and DataLoader
2. Define the model using nn.Module
3. Select a loss function and optimizer
4. Iterate over data, compute output and loss, and update weights

Sample Dataset

import torch
from torch.utils.data import DataLoader, TensorDataset

# Dummy data
X = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[2.0], [4.0], [6.0], [8.0]])

# Wrap in DataLoader
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=2, shuffle=True)

Define a Simple Model

import torch.nn as nn

class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = nn.Linear(1, 1)

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

model = LinearModel()

Loss Function and Optimizer

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

Training Loop

epochs = 100
for epoch in range(epochs):
    for batch_X, batch_y in loader:
        # Forward pass
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        
        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

Saving and Loading Model

# Save model
torch.save(model.state_dict(), 'linear_model.pth')

# Load model
model2 = LinearModel()
model2.load_state_dict(torch.load('linear_model.pth'))

Adding Validation During Training

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset, random_split

X = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]])
y = torch.tensor([[2.0], [4.0], [6.0], [8.0], [10.0]])

dataset = TensorDataset(X, y)
train_set, val_set = random_split(dataset, [4, 1])

train_loader = DataLoader(train_set, batch_size=2, shuffle=True)
val_loader = DataLoader(val_set, batch_size=1)

class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = nn.Linear(1, 1)

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

model = LinearModel()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

epochs = 50
for epoch in range(epochs):
    # Training phase
    model.train()
    for batch_X, batch_y in train_loader:
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Validation phase
    model.eval()
    with torch.no_grad():
        for val_X, val_y in val_loader:
            val_output = model(val_X)
            val_loss = criterion(val_output, val_y)

    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Train Loss: {loss.item():.4f}, Val Loss: {val_loss.item():.4f}")

• Uses random_split() to divide data into training and validation sets.
• model.train() and model.eval() set the mode for training and inference, respectively.
• torch.no_grad() is used during validation to reduce memory usage.

Use Case: Batch Training with Multiple Features

X = torch.tensor([
    [1.0, 2.0],
    [2.0, 3.0],
    [3.0, 4.0],
    [4.0, 5.0]
])
y = torch.tensor([[5.0], [7.0], [9.0], [11.0]])

dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=2, shuffle=True)

class TwoFeatureLinear(nn.Module):
    def __init__(self):
        super(TwoFeatureLinear, self).__init__()
        self.linear = nn.Linear(2, 1)

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

model = TwoFeatureLinear()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

for epoch in range(30):
    for batch_X, batch_y in loader:
        pred = model(batch_X)
        loss = criterion(pred, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

print("Final Weights:", model.linear.weight)
print("Final Bias:", model.linear.bias)

• This example demonstrates training a model with two input features using the Adam optimizer.
• At the end of training, we output learned weights and bias to inspect model parameters.

Saving and Loading Models

# Save the model weights only
torch.save(model.state_dict(), 'model_weights.pth')

# To load later
loaded_model = TwoFeatureLinear()
loaded_model.load_state_dict(torch.load('model_weights.pth'))
loaded_model.eval()

• Always call .eval() after loading if you're performing inference.
• state_dict() contains all learnable weights and biases in the model.

Conclusion

With PyTorch, training a model involves a clear and flexible loop for computing predictions, evaluating loss, and updating weights. DataLoader efficiently handles batching and shuffling, making training both simple and scalable.

Subhendu Mohapatra

Author & Instructor at plus2net

I write and maintain practical tutorials on Python, PHP, SQL, JavaScript, HTML, jQuery, and web development at plus2net. The tutorials focus on clear explanations, working examples, and code that readers can test and adapt while learning.

About the author   Watch on YouTube

Connect: