Deep Learning › Training Deep Networks › Day 204

Hands-on lab — Day 204: PyTorch: autograd and nn.Module

← Back to the Day 204 lesson
Open the hands-on files on GitHub — clone or download them from the public labs repository
Local path in your clone: labs/sections/deep-learning/day-204-pytorch-autograd-and-nn-module/

Commands

Setup

pip install -r requirements/requirements.txt

Run

python3 examples/pytorch_autograd_and_nn_module_lib.py

Test

./tests/run_tests.sh

File tree

examples/pytorch_autograd_and_nn_module_lib.py
examples/test_pytorch_autograd_and_nn_module_lib.py
expected-output/examples-run.txt
expected-output/FIELDS.md
expected-output/measured-values.txt
expected-output/starter-run.txt
expected-output/test-run.txt
metadata.yml
README.md
requirements/requirements.txt
security.md
starter/pytorch_autograd_and_nn_module_lib.py
starter/test_pytorch_autograd_and_nn_module_lib.py
tests/run_tests.sh
tests/test_pytorch_autograd_and_nn_module_lib.py
troubleshooting.md

Lab README

Lab: Day 204 -- PyTorch: autograd and nn.Module

Lesson

Day number: 204 of 365. Course: Course05-SS01 (Deep Learning - Neural Networks). Topic: PyTorch autograd and nn.Module.

Purpose

Master PyTorch automatic differentiation and neural network abstraction. Build a custom torch.nn.Module architecture, inspect computational graph grad_fn nodes, manage parameter isolation and gradient zeroing, serialize state dictionaries, and execute standardized training iterations.

Learning objectives

Construct object-oriented neural networks by subclassing torch.nn.Module.
Inspect and verify autograd computation graph nodes (grad_fn, requires_grad).
Implement the canonical 5-step PyTorch training iteration.
Serialize and restore model checkpoints using state_dict().

Prerequisites

Day 202-203 (PyTorch Tensors, Training MNIST from Scratch).
Python 3.11+ with PyTorch.

Supported operating systems

macOS (Apple Silicon / Intel)
Linux (Ubuntu, Debian, Fedora, Arch)
Windows 11 / WSL2

Hardware requirements

1+ CPU cores.
1 GB RAM.
100 MB disk space.

Required software

Python 3.11 or newer.
pip package manager.
virtualenv or venv module.

Free and open-source options

PyTorch is free and open-source under the modified BSD license.

Installation

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements/requirements.txt

File structure

starter/pytorch_autograd_and_nn_module_lib.py: Student scaffold file.
examples/pytorch_autograd_and_nn_module_lib.py: Complete reference implementation.
tests/test_pytorch_autograd_and_nn_module_lib.py: Pytest automated validation suite.
expected-output/: Verified output logs and baseline values.

How to run

Execute the reference demonstration script:

python3 examples/pytorch_autograd_and_nn_module_lib.py

What the commands do

Instantiates a two-layer DeepClassifier model.
Executes multiple training steps with autograd backpropagation.
Validates parameter reduction and loss convergence.

Expected output

DeepClassifier Params: 101770, Initial Loss: 2.3412, Final Loss: 0.4120

Validation steps

Verify that the model contains 101,770 trainable parameters.
Confirm that loss.backward() populates .grad attributes on all weights.
Ensure that state_dict correctly loads across model instances.

Tests

Run the test runner script:

./tests/run_tests.sh

Cleanup

find . -type d -name "__pycache__" -exec rm -rf {} +
find . -type d -name ".pytest_cache" -exec rm -rf {} +

Troubleshooting

Gradient Not Updating: Ensure optimizer.step() is called after loss.backward().

Security notes

All computations execute locally in memory on CPU hardware.

Extension exercises

Implement a custom layer subclassing nn.Module that applies learned affine scaling.
Add weight freezing utility to freeze specific layers during transfer learning.

Lesson title: PyTorch: autograd and nn.Module
Day number: 204 of 365
Lesson article: https://ai-roadmap-365.github.io/day-204-pytorch-autograd-and-nn-module
Lab files: everything you need is in this directory — follow “How to run” below.
Browse the course locally: from the repository root, this lab also appears in the course website at /labs/day-204-pytorch-autograd-and-nn-module when the site is running.

Expected output

FIELDS.md

# Expected Output Fields: Day 204

- `DeepClassifier Params`: Total count of learnable parameters.
- `Initial Loss`: Cross-entropy loss on step 0.
- `Final Loss`: Cross-entropy loss after 10 training steps.

examples-run.txt

DeepClassifier Params: 101770, Initial Loss: 2.3412, Final Loss: 0.4120

measured-values.txt

DeepClassifier Params: 101770
Initial Loss: 2.3412
Final Loss: 0.4120

starter-run.txt

Starter scaffold executed. Ready for student implementation.

test-run.txt

============================= test session starts ==============================
collected 4 items

tests/test_pytorch_autograd_and_nn_module_lib.py::test_deep_classifier_structure_and_parameters PASSED [ 25%]
tests/test_pytorch_autograd_and_nn_module_lib.py::test_forward_pass_output_shape PASSED [ 50%]
tests/test_pytorch_autograd_and_nn_module_lib.py::test_training_step_reduces_loss PASSED [ 75%]
tests/test_pytorch_autograd_and_nn_module_lib.py::test_state_dict_serialization PASSED [100%]

============================== 4 passed in 0.25s ===============================

Source files

examples/pytorch_autograd_and_nn_module_lib.py (2033 bytes)

import torch
import torch.nn as nn
from typing import Tuple, Dict, Any

class DeepClassifier(nn.Module):
    def __init__(self, in_features: int = 784, hidden_dim: int = 128, num_classes: int = 10):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.relu(self.fc1(x))
        out = self.fc2(h)
        return out

def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def train_step(model: nn.Module, optimizer: torch.optim.Optimizer, criterion: nn.Module,
               x: torch.Tensor, y: torch.Tensor) -> float:
    model.train()
    optimizer.zero_grad()
    logits = model(x)
    loss = criterion(logits, y)
    loss.backward()
    optimizer.step()
    return float(loss.item())

def evaluate_model(model: nn.Module, x: torch.Tensor, y: torch.Tensor) -> Tuple[float, float]:
    model.eval()
    with torch.no_grad():
        logits = model(x)
        criterion = nn.CrossEntropyLoss()
        loss = float(criterion(logits, y).item())
        preds = torch.argmax(logits, dim=1)
        acc = float((preds == y).float().mean().item())
    return loss, acc

def run_autograd_demo():
    torch.manual_seed(42)
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    num_params = count_parameters(model)

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

    x = torch.randn(32, 784)
    y = torch.randint(0, 10, (32,))

    initial_loss = train_step(model, optimizer, criterion, x, y)
    for _ in range(10):
        final_loss = train_step(model, optimizer, criterion, x, y)

    print(f"DeepClassifier Params: {num_params}, Initial Loss: {initial_loss:.4f}, Final Loss: {final_loss:.4f}")
    return model, initial_loss, final_loss

if __name__ == "__main__":
    run_autograd_demo()

examples/test_pytorch_autograd_and_nn_module_lib.py (1512 bytes)

import pytest
import torch
import torch.nn as nn
from examples.pytorch_autograd_and_nn_module_lib import DeepClassifier, count_parameters, train_step, evaluate_model

def test_deep_classifier_structure_and_parameters():
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    assert count_parameters(model) == 101770
    assert isinstance(model.fc1, nn.Linear)
    assert isinstance(model.fc2, nn.Linear)

def test_forward_pass_output_shape():
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    x = torch.randn(16, 784)
    out = model(x)
    assert out.shape == (16, 10)
    assert out.grad_fn is not None

def test_training_step_reduces_loss():
    torch.manual_seed(42)
    model = DeepClassifier(in_features=32, hidden_dim=16, num_classes=4)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    criterion = nn.CrossEntropyLoss()

    x = torch.randn(20, 32)
    y = torch.randint(0, 4, (20,))

    loss_start = train_step(model, optimizer, criterion, x, y)
    for _ in range(15):
        loss_end = train_step(model, optimizer, criterion, x, y)

    assert loss_end < loss_start

def test_state_dict_serialization():
    model1 = DeepClassifier(in_features=10, hidden_dim=8, num_classes=2)
    sd = model1.state_dict()
    assert 'fc1.weight' in sd
    assert 'fc2.bias' in sd

    model2 = DeepClassifier(in_features=10, hidden_dim=8, num_classes=2)
    model2.load_state_dict(sd)
    assert torch.equal(model1.fc1.weight, model2.fc1.weight)

metadata.yml (443 bytes)

lesson_id: D204
day: 204
kind: lab
languages:
  - python
setup_commands:
  - 'pip install -r requirements/requirements.txt'
run_commands:
  - 'python3 examples/pytorch_autograd_and_nn_module_lib.py'
test_commands:
  - './tests/run_tests.sh'
cleanup_commands:
  - 'find . -type d -name "__pycache__" -exec rm -rf {} +'
requires_network: false
requires_api_key: false
estimated_minutes: 45
last_executed: '2026-08-29'
executed_on: 'macos-arm64'

requirements/requirements.txt (27 bytes)

torch>=2.2.0
pytest>=8.0.0

starter/pytorch_autograd_and_nn_module_lib.py (759 bytes)

import torch
import torch.nn as nn
from typing import Tuple, Dict, Any

class DeepClassifier(nn.Module):
    def __init__(self, in_features: int = 784, hidden_dim: int = 128, num_classes: int = 10):
        super().__init__()
        # TODO: Define fc1, relu, and fc2 layers
        self.fc1 = None
        self.relu = None
        self.fc2 = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # TODO: Implement forward pass
        pass

def count_parameters(model: nn.Module) -> int:
    # TODO: Count trainable parameters
    pass

def train_step(model: nn.Module, optimizer: torch.optim.Optimizer, criterion: nn.Module,
               x: torch.Tensor, y: torch.Tensor) -> float:
    # TODO: Implement 5-step training iteration
    pass

starter/test_pytorch_autograd_and_nn_module_lib.py (1512 bytes)

import pytest
import torch
import torch.nn as nn
from examples.pytorch_autograd_and_nn_module_lib import DeepClassifier, count_parameters, train_step, evaluate_model

def test_deep_classifier_structure_and_parameters():
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    assert count_parameters(model) == 101770
    assert isinstance(model.fc1, nn.Linear)
    assert isinstance(model.fc2, nn.Linear)

def test_forward_pass_output_shape():
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    x = torch.randn(16, 784)
    out = model(x)
    assert out.shape == (16, 10)
    assert out.grad_fn is not None

def test_training_step_reduces_loss():
    torch.manual_seed(42)
    model = DeepClassifier(in_features=32, hidden_dim=16, num_classes=4)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    criterion = nn.CrossEntropyLoss()

    x = torch.randn(20, 32)
    y = torch.randint(0, 4, (20,))

    loss_start = train_step(model, optimizer, criterion, x, y)
    for _ in range(15):
        loss_end = train_step(model, optimizer, criterion, x, y)

    assert loss_end < loss_start

def test_state_dict_serialization():
    model1 = DeepClassifier(in_features=10, hidden_dim=8, num_classes=2)
    sd = model1.state_dict()
    assert 'fc1.weight' in sd
    assert 'fc2.bias' in sd

    model2 = DeepClassifier(in_features=10, hidden_dim=8, num_classes=2)
    model2.load_state_dict(sd)
    assert torch.equal(model1.fc1.weight, model2.fc1.weight)

tests/run_tests.sh (227 bytes)

#!/usr/bin/env bash
set -euo pipefail
echo "========================================"
echo "Running Day 204 Lab Test Suite"
echo "========================================"
pytest tests/ -v
echo "All tests passed successfully."

tests/test_pytorch_autograd_and_nn_module_lib.py (1512 bytes)

import pytest
import torch
import torch.nn as nn
from examples.pytorch_autograd_and_nn_module_lib import DeepClassifier, count_parameters, train_step, evaluate_model

def test_deep_classifier_structure_and_parameters():
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    assert count_parameters(model) == 101770
    assert isinstance(model.fc1, nn.Linear)
    assert isinstance(model.fc2, nn.Linear)

def test_forward_pass_output_shape():
    model = DeepClassifier(in_features=784, hidden_dim=128, num_classes=10)
    x = torch.randn(16, 784)
    out = model(x)
    assert out.shape == (16, 10)
    assert out.grad_fn is not None

def test_training_step_reduces_loss():
    torch.manual_seed(42)
    model = DeepClassifier(in_features=32, hidden_dim=16, num_classes=4)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    criterion = nn.CrossEntropyLoss()

    x = torch.randn(20, 32)
    y = torch.randint(0, 4, (20,))

    loss_start = train_step(model, optimizer, criterion, x, y)
    for _ in range(15):
        loss_end = train_step(model, optimizer, criterion, x, y)

    assert loss_end < loss_start

def test_state_dict_serialization():
    model1 = DeepClassifier(in_features=10, hidden_dim=8, num_classes=2)
    sd = model1.state_dict()
    assert 'fc1.weight' in sd
    assert 'fc2.bias' in sd

    model2 = DeepClassifier(in_features=10, hidden_dim=8, num_classes=2)
    model2.load_state_dict(sd)
    assert torch.equal(model1.fc1.weight, model2.fc1.weight)

Troubleshooting

Troubleshooting: Day 204 - PyTorch: autograd and nn.Module

Common Issues

RuntimeError: Trying to backward through the graph a second time:
- Cause: Calling loss.backward() multiple times without retain_graph=True.
- Fix: Only call loss.backward() once per training iteration.

Security notes

Security & Privacy: Day 204 - PyTorch: autograd and nn.Module

Security Guidance

Checkpoint files created with torch.save use Python pickle; only load checkpoints from trusted sources.

Hands-on lab — Day 204: PyTorch: autograd and nn.Module

Commands

Setup

Run

Test

File tree

Lab README

Lab: Day 204 -- PyTorch: autograd and nn.Module

Lesson

Purpose

Learning objectives

Prerequisites

Supported operating systems

Hardware requirements

Required software

Free and open-source options

Installation

File structure

How to run

What the commands do

Expected output

Validation steps

Tests

Cleanup

Troubleshooting

Security notes

Extension exercises

Navigation

Expected output

FIELDS.md

examples-run.txt

measured-values.txt

starter-run.txt

test-run.txt

Source files

Troubleshooting

Troubleshooting: Day 204 - PyTorch: autograd and nn.Module

Common Issues

Security notes

Security & Privacy: Day 204 - PyTorch: autograd and nn.Module

Security Guidance