Capture mode in practice

This notebook uses DALI Dynamic to read and normalize MNIST images for a small PyTorch classifier, then measures the same training loop with and without capture mode.

It requires PyTorch and DALI_extra. Set DALI_EXTRA_PATH to the DALI_extra directory.

Note

See execution model to understand how capture mode works and capture rules for the rules that decide whether an operator is captured at all.

Dataset and model

DALI_extra stores MNIST as JPEG images in a Caffe2-compatible LMDB. The model is deliberately small so that it demonstrates a real training loop without obscuring the data-loading example.

import os
import statistics
from pathlib import Path
from timeit import default_timer as timer

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

import nvidia.dali.experimental.dynamic as ndd
import nvidia.dali.types as types

DATA_PATH = Path(os.environ["DALI_EXTRA_PATH"]) / "db" / "MNIST" / "training"
BATCH_SIZE = 64


class MNISTClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, images):
        images = torch.flatten(images, start_dim=1)
        return self.fc2(F.relu(self.fc1(images)))

Preprocessing

preprocess is ordinary dynamic mode code. It decodes grayscale images, rotates them randomly, converts them to the CHW layout expected by PyTorch, and normalizes them with the standard MNIST mean and standard deviation.

def preprocess(
    jpegs: ndd.Batch,
    labels: ndd.Batch,
    rng: ndd.random.RNG,
) -> tuple[torch.Tensor, torch.Tensor]:
    images = ndd.decoders.image(jpegs, device="gpu", output_type=types.GRAY)
    angles = ndd.random.uniform(
        batch_size=BATCH_SIZE,
        range=[-7.5, 7.5],
        rng=rng,
    )
    images = ndd.rotate(images, angle=angles, keep_size=True, fill_value=0)
    images = ndd.crop_mirror_normalize(
        images,
        dtype=ndd.float32,
        output_layout="CHW",
        mean=[0.1307 * 255],
        std=[0.3081 * 255],
    )

    labels = labels.gpu()
    labels = ndd.squeeze(labels, axes=0)
    labels = ndd.cast(labels, dtype=ndd.int64)

    return images.torch(), labels.torch()

Training and benchmarking

The reader and model are fresh for each benchmark. The warmup epoch initializes DALI and the model. In capture mode, it also covers tracing and pipeline construction.

The only capture-mode-specific code is capture=capture_mode in next_epoch. The preprocessing and training body is otherwise identical.

def train_epoch(reader, rng, model, optimizer, *, capture_mode):
    loss = None
    samples = 0
    for jpegs, labels in reader.next_epoch(
        batch_size=BATCH_SIZE, capture=capture_mode
    ):
        images, targets = preprocess(jpegs, labels, rng)

        optimizer.zero_grad(set_to_none=True)
        loss = F.cross_entropy(model(images), targets)
        loss.backward()
        optimizer.step()

        samples += targets.numel()

    torch.cuda.synchronize()
    return loss.item(), samples


def benchmark(capture_mode: bool) -> tuple[float, float]:
    reader = ndd.readers.Caffe2(path=DATA_PATH, random_shuffle=False)
    rng = ndd.random.RNG(0)
    torch.manual_seed(0)

    model = MNISTClassifier().cuda()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    throughputs = []
    start = None
    for _ in range(10):
        loss, samples = train_epoch(
            reader,
            rng,
            model,
            optimizer,
            capture_mode=capture_mode,
        )
        if start is not None:
            throughputs.append(samples / (timer() - start))
        start = timer()

    return statistics.mean(throughputs), loss

Let’s run the benchmark in eager and capture mode and observe the difference.

eager_throughput, eager_loss = benchmark(capture_mode=False)
capture_throughput, capture_loss = benchmark(capture_mode=True)

print(
    f"eager:    {eager_throughput:,.0f} samples/s, final loss {eager_loss:.4f}"
)
print(
    f"capture:  {capture_throughput:,.0f} samples/s, final loss {capture_loss:.4f}"
)
print(
    f"speedup:  {(capture_throughput - eager_throughput) / eager_throughput:.2%}"
)

eager:    25,713 samples/s, final loss 0.0941
capture:  41,392 samples/s, final loss 0.0941
speedup:  60.98%

Capture mode removes most per-operator Python overhead and prefetches later batches while PyTorch trains on the current one. The exact speedup depends on the machine, the workload, and the balance between preprocessing and model work.