Testing PyTorch Applications: Unit Tests to Integration

Testing PyTorch Applications: Unit Tests to Integration

PyTorch has become the de facto framework for deep learning research and production, but testing machine learning code remains an afterthought for many teams. Unlike traditional software, ML code combines deterministic logic with stochastic behavior, making it tricky to test reliably. This tutorial walks you through a complete testing strategy for PyTorch applications — from granular unit tests on tensor operations to full integration tests that validate training pipelines end to end.

Why Testing PyTorch Code Matters

Machine learning bugs are silent. A wrong tensor reshape, a flipped dimension, or a detached gradient will not raise an exception — it will simply produce a slightly worse model. Without tests, these defects slip into production and degrade metrics slowly enough that nobody notices until the damage is done. A solid test suite catches shape mismatches early, documents expected behavior, and gives you the confidence to refactor model architectures without fear of breaking the training loop.

Setting Up the Testing Environment

Most PyTorch projects use pytest because of its clean assertion syntax and powerful fixtures. Install it alongside a few helper libraries:

pip install pytest pytest-cov torch torchvision numpy

Create a project structure that separates source code from tests:

my_model/
├── my_model/
│   ├── __init__.py
│   ├── data.py
│   ├── model.py
│   └── train.py
├── tests/
│   ├── __init__.py
│   ├── conftest.py
│   ├── test_data.py
│   ├── test_model.py
│   └── test_train.py
└── pytest.ini

A minimal pytest.ini keeps configuration explicit:

[pytest]
testpaths = tests
addopts = -v --tb=short
python_files = test_*.py

Unit Testing Tensor Operations and Data Pipelines

Unit tests focus on the smallest pieces of logic in isolation. For PyTorch, that means testing custom datasets, transforms, collate functions, and any tensor manipulation you write yourself. The goal is speed: these tests should run in milliseconds and never touch a GPU.

Testing a Custom Dataset

Suppose you have a dataset class that loads images and applies augmentations. The test should verify shapes, dtypes, and that indexing returns the expected structure.

# my_model/data.py
import torch
from torch.utils.data import Dataset

class SyntheticImageDataset(Dataset):
    def __init__(self, num_samples=100, image_size=(3, 32, 32), num_classes=10):
        self.num_samples = num_samples
        self.image_size = image_size
        self.num_classes = num_classes
        self.images = torch.randn(num_samples, *image_size)
        self.labels = torch.randint(0, num_classes, (num_samples,))

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]

# tests/test_data.py
import pytest
import torch
from my_model.data import SyntheticImageDataset

def test_dataset_length():
    dataset = SyntheticImageDataset(num_samples=50)
    assert len(dataset) == 50

def test_dataset_item_shapes():
    dataset = SyntheticImageDataset(num_samples=10, image_size=(3, 32, 32))
    image, label = dataset[0]
    assert image.shape == (3, 32, 32)
    assert image.dtype == torch.float32
    assert isinstance(label.item(), int)

def test_dataset_index_out_of_range():
    dataset = SyntheticImageDataset(num_samples=5)
    with pytest.raises(IndexError):
        _ = dataset[10]

def test_dataset_batch_consistency():
    dataset = SyntheticImageDataset(num_samples=20)
    loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=False)
    images, labels = next(iter(loader))
    assert images.shape == (4, 3, 32, 32)
    assert labels.shape == (4,)

Testing Custom Transforms

Transforms are pure functions, which makes them ideal unit test targets. Verify both the output shape and that the transform is invertible where applicable.

# my_model/data.py
class NormalizeImage:
    def __init__(self, mean, std):
        self.mean = torch.tensor(mean).view(-1, 1, 1)
        self.std = torch.tensor(std).view(-1, 1, 1)

    def __call__(self, image):
        return (image - self.mean) / self.std

# tests/test_data.py
def test_normalize_transform():
    transform = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    image = torch.ones(3, 8, 8)
    out = transform(image)
    assert out.shape == image.shape
    assert torch.allclose(out, torch.zeros(3, 8, 8), atol=1e-6)

Unit Testing Model Architecture

Model tests verify that your network accepts the right input shapes, produces the right output shapes, and that parameters update during a backward pass. These tests use tiny random tensors rather than real data so they run instantly.

Shape and Forward Pass Tests

# my_model/model.py
import torch.nn as nn

class TinyCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Linear(32 * 8 * 8, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# tests/test_model.py
import pytest
import torch
from my_model.model import TinyCNN

def test_forward_output_shape():
    model = TinyCNN(num_classes=10)
    x = torch.randn(4, 3, 32, 32)
    out = model(x)
    assert out.shape == (4, 10)

def test_forward_single_sample():
    model = TinyCNN(num_classes=5)
    x = torch.randn(1, 3, 32, 32)
    out = model(x)
    assert out.shape == (1, 5)

def test_forward_wrong_input_shape_raises():
    model = TinyCNN(num_classes=10)
    x = torch.randn(4, 1, 32, 32)  # wrong channel count
    with pytest.raises(RuntimeError):
        model(x)

def test_parameters_update_on_backward():
    model = TinyCNN(num_classes=10)
    x = torch.randn(2, 3, 32, 32)
    target = torch.tensor([3, 7])
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    optimizer.zero_grad()
    out = model(x)
    loss = criterion(out, target)
    loss.backward()

    # At least one parameter should have a non-None gradient
    has_grad = any(p.grad is not None and p.grad.abs().sum() > 0
                   for p in model.parameters())
    assert has_grad

Testing Parameter Count and Initialization

It is surprisingly easy to accidentally double a layer or forget to initialize a submodule. A quick test on parameter count guards against regressions when you refactor.

# tests/test_model.py
def test_parameter_count_in_expected_range():
    model = TinyCNN(num_classes=10)
    num_params = sum(p.numel() for p in model.parameters())
    # Roughly 200k parameters for this architecture
    assert 100_000 < num_params < 500_000

def test_no_nan_parameters_after_init():
    model = TinyCNN(num_classes=10)
    for name, param in model.named_parameters():
        assert not torch.isnan(param).any(), f"NaN in {name}"

Testing Training Logic

The training step is where most subtle bugs live: forgetting optimizer.zero_grad(), detaching tensors that should carry gradients, or computing loss on the wrong device. Test the training step as a function that takes a batch and returns a loss, so you can assert it decreases over a few iterations on synthetic data.

Extracting a Testable Training Step

# my_model/train.py
import torch
import torch.nn as nn

def train_step(model, batch, criterion, optimizer, device="cpu"):
    model.train()
    images, labels = batch
    images = images.to(device)
    labels = labels.to(device)

    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    return loss.item()

def evaluate(model, loader, criterion, device="cpu"):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            total_loss += criterion(outputs, labels).item() * images.size(0)
            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)
    return total_loss / total, correct / total

# tests/test_train.py
import pytest
import torch
import torch.nn as nn
from my_model.model import TinyCNN
from my_model.data import SyntheticImageDataset
from my_model.train import train_step, evaluate

@pytest.fixture
def small_setup():
    torch.manual_seed(42)
    model = TinyCNN(num_classes=10)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
    dataset = SyntheticImageDataset(num_samples=32, num_classes=10)
    return model, criterion, optimizer, dataset

def test_train_step_returns_float(small_setup):
    model, criterion, optimizer, dataset = small_setup
    batch = (dataset.images[:4], dataset.labels[:4])
    loss = train_step(model, batch, criterion, optimizer)
    assert isinstance(loss, float)
    assert loss >= 0

def test_loss_decreases_over_steps(small_setup):
    model, criterion, optimizer, dataset = small_setup
    batch = (dataset.images[:8], dataset.labels[:8])
    losses = [train_step(model, batch, criterion, optimizer) for _ in range(20)]
    assert losses[-1] < losses[0]

def test_gradients_are_cleared_after_step(small_setup):
    model, criterion, optimizer, dataset = small_setup
    batch = (dataset.images[:4], dataset.labels[:4])
    train_step(model, batch, criterion, optimizer)
    # After zero_grad at the start of next step, grads should reset
    # Verify by checking that a second step does not accumulate
    grad_before = model.classifier.weight.grad.clone()
    train_step(model, batch, criterion, optimizer)
    grad_after = model.classifier.weight.grad
    # Gradients should differ (not simply doubled)
    assert not torch.allclose(grad_before, grad_after)

def test_evaluate_returns_metrics(small_setup):
    model, criterion, optimizer, dataset = small_setup
    loader = torch.utils.data.DataLoader(dataset, batch_size=8)
    avg_loss, accuracy = evaluate(model, loader, criterion)
    assert isinstance(avg_loss, float)
    assert 0.0 <= accuracy <= 1.0

Using Fixtures and Parametrization Effectively

Pytest fixtures let you share expensive setup across tests, and parametrization lets you run the same test against multiple configurations. For PyTorch, a common pattern is parametrizing over batch sizes, input dimensions, and devices.

# tests/conftest.py
import pytest
import torch

@pytest.fixture(params=[1, 4, 16])
def batch_size(request):
    return request.param

@pytest.fixture(params=[(3, 32, 32), (3, 64, 64)])
def image_size(request):
    return request.param

@pytest.fixture
def random_batch(batch_size, image_size):
    images = torch.randn(batch_size, *image_size)
    labels = torch.randint(0, 10, (batch_size,))
    return images, labels

# tests/test_model.py
def test_model_handles_various_batches(random_batch):
    model = TinyCNN(num_classes=10)
    images, labels = random_batch
    # Adjust model for different input sizes by using adaptive pooling
    # For this example, we only test the 32x32 case
    if images.shape[-1] == 32:
        out = model(images)
        assert out.shape == (images.size(0), 10)

Integration Testing the Full Pipeline

Integration tests validate that all components work together: data loading, model forward and backward passes, checkpointing, and evaluation. These tests run longer than unit tests but catch wiring mistakes that unit tests miss. The trick is to use a tiny dataset and a minimal number of epochs so the test still completes in a few seconds.

End-to-End Training Integration Test

# tests/test_train.py
import tempfile
import os
import torch
import torch.nn as nn
from my_model.model import TinyCNN
from my_model.data import SyntheticImageDataset
from my_model.train import train_step, evaluate

def test_full_training_loop_improves_accuracy():
    torch.manual_seed(0)
    model = TinyCNN(num_classes=5)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    # Use a small, learnable synthetic dataset: random images with
    # labels derived from a simple rule so the model can actually learn
    num_samples = 64
    images = torch.randn(num_samples, 3, 32, 32)
    labels = (images.mean(dim=(1, 2, 3)) > 0).long()
    dataset = torch.utils.data.TensorDataset(images, labels)
    loader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)

    initial_loss, initial_acc = evaluate(model, loader, criterion)

    for epoch in range(15):
        for batch_images, batch_labels in loader:
            train_step(model, (batch_images, batch_labels), criterion, optimizer)

    final_loss, final_acc = evaluate(model, loader, criterion)

    assert final_acc > initial_acc
    assert final_loss < initial_loss
    assert final_acc > 0.6  # should learn the simple rule reasonably well

def test_checkpoint_save_and_load():
    torch.manual_seed(123)
    model = TinyCNN(num_classes=10)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    # Run one step to populate optimizer state
    x = torch.randn(4, 3, 32, 32)
    y = torch.tensor([0, 1, 2, 3])
    train_step(model, (x, y), criterion, optimizer)

    with tempfile.TemporaryDirectory() as tmpdir:
        ckpt_path = os.path.join(tmpdir, "model.pt")
        torch.save({
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
        }, ckpt_path)

        # Load into a fresh model
        new_model = TinyCNN(num_classes=10)
        new_optimizer = torch.optim.SGD(new_model.parameters(), lr=0.01)
        ckpt = torch.load(ckpt_path)
        new_model.load_state_dict(ckpt["model_state_dict"])
        new_optimizer.load_state_dict(ckpt["optimizer_state_dict"])

        # Both models should produce identical outputs
        new_model.eval()
        model.eval()
        with torch.no_grad():
            out1 = model(x)
            out2 = new_model(x)
        assert torch.allclose(out1, out2, atol=1e-6)

Testing Device Placement

If your code supports GPU, guard the test with a skip marker so it only runs when CUDA is available. This keeps the suite portable across CI environments.

# tests/test_train.py
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_training_step_runs_on_gpu():
    device = "cuda"
    model = TinyCNN(num_classes=10).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    images = torch.randn(4, 3, 32, 32, device=device)
    labels = torch.tensor([0, 1, 2, 3], device=device)
    loss = train_step(model, (images, labels), criterion, optimizer, device=device)
    assert isinstance(loss, float)
    # Verify model parameters are still on GPU
    assert next(model.parameters()).is_cuda

Testing Numerical Stability and Determinism

Stochastic behavior is the enemy of reproducible tests. Set seeds at the start of every test that involves randomness, and consider testing that your model produces deterministic outputs in eval mode.

# tests/test_model.py
def test_eval_mode_is_deterministic():
    torch.manual_seed(99)
    model = TinyCNN(num_classes=10)
    model.eval()
    x = torch.randn(2, 3, 32, 32)
    with torch.no_grad():
        out1 = model(x)
        out2 = model(x)
    assert torch.allclose(out1, out2)

def test_batchnorm_uses_running_stats_in_eval():
    model = nn.Sequential(
        nn.Conv2d(3, 8, 3, padding=1),
        nn.BatchNorm2d(8),
        nn.AdaptiveAvgPool2d(1),
        nn.Flatten(),
        nn.Linear(8, 10),
    )
    model.train()
    x = torch.randn(8, 3, 32, 32)
    model(x)  # update running stats
    model.eval()
    with torch.no_grad():
        single = model(x[:1])
        batched = model(x)[:1]
    # In eval mode, output should not depend on batch composition
    assert torch.allclose(single, batched, atol=1e-5)

Best Practices for PyTorch Test Suites

• Keep tests fast. Use tiny synthetic tensors and single-batch training steps. Save real datasets for a separate, slower integration test suite.
• Always set random seeds. Call torch.manual_seed() at the start of any test that touches randomness, including model initialization.
• Test shapes before values. Most PyTorch bugs are shape errors. Asserting output shapes catches them immediately.
• Separate unit and integration tests. Use pytest markers (@pytest.mark.slow) so you can run the fast suite on every save and the full suite before merging.
• Avoid testing on GPU in CI unless necessary. GPU tests are slower and harder to reproduce. Run them in a nightly job instead of on every commit.
• Test the contract, not the implementation. Assert that loss decreases and accuracy improves, not that specific weights have specific values. Implementation tests break every time you refactor.
• Use torch.allclose with explicit tolerances. Floating point comparisons with == will fail unpredictably across platforms.
• Test edge cases explicitly. Batch size of one, empty sequences, and extreme input values often expose bugs that normal batches hide.
• Snapshot model outputs sparingly. Snapshot tests are brittle when architecture changes. Prefer behavioral assertions like "loss decreases" or "accuracy above threshold."

Conclusion

Testing PyTorch applications requires a shift in mindset from traditional software testing. You are not just verifying that code runs without errors — you are verifying that tensors flow through the graph correctly, that gradients propagate where they should, and that the training loop actually learns. By layering fast unit tests for shapes and transforms, focused tests for the training step, and integration tests for the full pipeline, you build a safety net that catches the silent failures unique to machine learning. Start small with shape assertions and a single loss-decrease test, then expand the suite as your model grows. The investment pays off the first time you refactor an architecture and the test suite tells you exactly what broke before your users do.