Computer vision & ResNet Part 6

Part 6 · Practice

ResNet from scratch, and the experiment worth running

Working PyTorch for the blocks, the full architectures, and a reproduction of the degradation result that motivated the whole paper.

Working PyTorch for everything described in ResNet, plus a reproduction of the degradation result from The challenge. The degradation experiment is the one I would actually run: it is small enough to finish on a laptop and it demonstrates the failure the paper was written to fix.

Everything here targets PyTorch 2.x and runs on CUDA, MPS, or CPU.

The residual block

FigureWhat the code below builds

Loading figure. It needs JavaScript; the surrounding text stands on its own without it.

The basic block maps to BasicBlock and the bottleneck to Bottleneck in the listings that follow.

The entire idea is four lines of forward. Note the ordering: the shortcut is added before the final ReLU, not after, so the identity path stays linear through the block.

pythonresnet/blocks.py
import torch
import torch.nn as nn
import torch.nn.functional as F


def conv3x3(cin, cout, stride=1):
    """3x3 convolution with padding. No bias: the BatchNorm that follows has one."""
    return nn.Conv2d(cin, cout, kernel_size=3, stride=stride, padding=1, bias=False)


class BasicBlock(nn.Module):
    """Two 3x3 convolutions. Used in ResNet-18 and ResNet-34."""

    expansion = 1

    def __init__(self, cin, planes, stride=1, downsample=None):
        super().__init__()
        self.conv1 = conv3x3(cin, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        identity = x

        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))          # F(x)

        if self.downsample is not None:
            identity = self.downsample(x)        # only when shape changes

        out = out + identity                     # F(x) + x
        return self.relu(out)


class Bottleneck(nn.Module):
    """1x1 reduce, 3x3, 1x1 expand. Used in ResNet-50, 101 and 152.

    `stride_on_3x3=False` reproduces the original paper, which strides in the
    first 1x1. `True` gives the widely used "v1.5" variant that strides in the
    3x3 instead, worth roughly half a point of top-1 accuracy because a strided
    1x1 convolution discards three quarters of its input outright.
    """

    expansion = 4

    def __init__(self, cin, planes, stride=1, downsample=None, stride_on_3x3=True):
        super().__init__()
        s1, s3 = (1, stride) if stride_on_3x3 else (stride, 1)
        self.conv1 = nn.Conv2d(cin, planes, 1, stride=s1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes, s3)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        identity = x

        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))

        if self.downsample is not None:
            identity = self.downsample(x)

        out = out + identity
        return self.relu(out)

The three shortcut options

The paper evaluates three ways to reconcile shapes when a stage halves resolution and changes channel count. Option A is the only one with no parameters at all, which is why the CIFAR experiments use it: the plain and residual networks then have identical parameter counts, so no one can attribute the difference to extra capacity.

pythonresnet/shortcuts.py
class ZeroPadShortcut(nn.Module):
    """Option A: subsample spatially, pad the new channels with zeros.

    Parameter-free, which makes plain-vs-residual comparisons exact.
    """

    def __init__(self, stride, extra_channels):
        super().__init__()
        self.stride = stride
        self.extra = extra_channels

    def forward(self, x):
        if self.stride > 1:
            x = x[:, :, ::self.stride, ::self.stride]
        lo = self.extra // 2
        return F.pad(x, (0, 0, 0, 0, lo, self.extra - lo))


def projection_shortcut(cin, cout, stride):
    """Option B: a 1x1 convolution, used only where dimensions change."""
    return nn.Sequential(
        nn.Conv2d(cin, cout, 1, stride=stride, bias=False),
        nn.BatchNorm2d(cout),
    )

The full ImageNet architectures

One class covers all five depths. The only thing that changes is the block type and the number of blocks per stage.

pythonresnet/model.py
class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False):
        super().__init__()
        self.cin = 64

        # stem: 224 -> 112 -> 56
        self.conv1 = nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)

        # four stages; channels double as resolution halves
        self.layer1 = self._stage(block, 64, layers[0], stride=1)   # 56x56
        self.layer2 = self._stage(block, 128, layers[1], stride=2)  # 28x28
        self.layer3 = self._stage(block, 256, layers[2], stride=2)  # 14x14
        self.layer4 = self._stage(block, 512, layers[3], stride=2)  # 7x7

        # global average pooling, not a big fully connected stack
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                # He initialization: the variance scaling derived for ReLU nets
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        if zero_init_residual:
            # Later refinement (Goyal et al. 2017): start every block as an exact
            # identity by zeroing the last BN's scale. Each block then begins life
            # as a no-op and has to earn its contribution.
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _stage(self, block, planes, blocks, stride):
        downsample = None
        cout = planes * block.expansion
        if stride != 1 or self.cin != cout:
            downsample = projection_shortcut(self.cin, cout, stride)  # option B

        layers = [block(self.cin, planes, stride, downsample)]
        self.cin = cout
        layers += [block(self.cin, planes) for _ in range(1, blocks)]
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
        x = self.layer4(self.layer3(self.layer2(self.layer1(x))))
        x = torch.flatten(self.avgpool(x), 1)
        return self.fc(x)


def resnet18():  return ResNet(BasicBlock, [2, 2, 2, 2])
def resnet34():  return ResNet(BasicBlock, [3, 4, 6, 3])
def resnet50():  return ResNet(Bottleneck, [3, 4, 6, 3])
def resnet101(): return ResNet(Bottleneck, [3, 4, 23, 3])
def resnet152(): return ResNet(Bottleneck, [3, 8, 36, 3])

Notice that ResNet-34 and ResNet-50 have the same [3, 4, 6, 3] block counts. Swapping two-layer basic blocks for three-layer bottleneck blocks is what takes it from 34 to 50 layers, and it barely changes the cost: 3.6 versus 3.8 billion multiply-adds.

Verifying against the paper

A self-contained counter, so you can check the numbers rather than trust them.

pythontools/complexity.py
import torch
import torch.nn as nn


@torch.no_grad()
def complexity(model, size=224):
    """Count parameters and forward multiply-adds (the paper's FLOP convention)."""
    macs = 0

    def conv_hook(module, inputs, output):
        nonlocal macs
        out_elems = output.numel() / output.shape[0]          # per image
        k = module.kernel_size[0] * module.kernel_size[1]
        macs += out_elems * k * (module.in_channels / module.groups)

    def linear_hook(module, inputs, output):
        nonlocal macs
        macs += module.in_features * module.out_features

    handles = []
    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            handles.append(m.register_forward_hook(conv_hook))
        elif isinstance(m, nn.Linear):
            handles.append(m.register_forward_hook(linear_hook))

    model.eval()
    model(torch.zeros(1, 3, size, size))
    for h in handles:
        h.remove()

    params = sum(p.numel() for p in model.parameters())
    return params, macs


if __name__ == "__main__":
    from resnet.model import resnet18, resnet34, resnet50, resnet101, resnet152

    print(f"{'model':<12}{'params':>10}{'G mult-adds':>14}{'paper':>10}")
    paper = {"resnet34": 3.6, "resnet50": 3.8, "resnet101": 7.6, "resnet152": 11.3}
    for fn in (resnet18, resnet34, resnet50, resnet101, resnet152):
        p, m = complexity(fn())
        ref = paper.get(fn.__name__, "")
        print(f"{fn.__name__:<12}{p/1e6:>9.1f}M{m/1e9:>14.2f}{ref:>10}")

Expected output, which matches the paper's quoted figures closely:

text
model           params   G mult-adds     paper
resnet18         11.7M          1.82
resnet34         21.8M          3.67       3.6
resnet50         25.6M          4.09       3.8
resnet101        44.5M          7.80       7.6
resnet152        60.2M         11.51      11.3

The small overshoot against the paper is the usual accounting difference over whether downsample convolutions and the classifier are included. It also explains the common discrepancy between "ResNet-50 is 3.8 GFLOPs" (the paper) and "4.1 GFLOPs" (most later work). Both mean multiply-adds. Double either figure if you want true floating point operation counts, since each multiply-add is two operations.

The identity-mapping thought experiment, in code

The paper's central argument is constructive, which means you can literally build the object it describes and confirm it behaves as claimed.

pythontools/identity_construction.py
import copy
import torch
import torch.nn as nn


class IdentityBlock(nn.Module):
    """A residual block wired to compute exactly f(x) = x.

    Zeroing the residual branch's output makes F(x) = 0, so the block returns
    relu(0 + x) = x for any non-negative input, which is what a post-ReLU
    activation always is.
    """

    def __init__(self, channels):
        super().__init__()
        self.body = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels, channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(channels),
        )
        nn.init.zeros_(self.body[-1].weight)   # BN scale -> 0
        nn.init.zeros_(self.body[-1].bias)

    def forward(self, x):
        return torch.relu(self.body(x) + x)


def deepen(shallow_net, channels, extra_blocks):
    """Build a deeper network guaranteed to compute the same function."""
    deeper = copy.deepcopy(shallow_net)
    deeper.extra = nn.Sequential(*[IdentityBlock(channels) for _ in range(extra_blocks)])
    return deeper


if __name__ == "__main__":
    torch.manual_seed(0)
    body = nn.Sequential(
        nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),
        nn.Conv2d(16, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),
    ).eval()

    x = torch.randn(4, 3, 32, 32)
    shallow_out = body(x)

    extra = nn.Sequential(*[IdentityBlock(16) for _ in range(20)]).eval()
    deep_out = extra(shallow_out)

    print("max |difference| after 20 extra blocks:",
          (deep_out - shallow_out).abs().max().item())
    # -> 0.0

Twenty extra layers, bit-identical output. That is the whole paradox in one number: this solution exists, it is trivial to write down, and stochastic gradient descent on a plain network cannot find it.

Reproducing degradation on CIFAR-10

This is the experiment I would spend an afternoon on. The paper's CIFAR networks are deliberately tiny, 0.27M to 1.7M parameters, and use option A shortcuts so the plain and residual versions have exactly the same parameter count.

Architecture: a 3x3 stem at 16 channels, then three stages of n blocks at 16, 32 and 64 channels, giving 6n + 2 layers. So n = 3 is ResNet-20 and n = 9 is ResNet-56.

pythonresnet/cifar.py
import torch
import torch.nn as nn

from resnet.blocks import BasicBlock, conv3x3
from resnet.shortcuts import ZeroPadShortcut


class PlainBlock(nn.Module):
    """BasicBlock with the shortcut deleted. Everything else is identical."""

    def __init__(self, cin, planes, stride=1, downsample=None):
        super().__init__()
        self.conv1 = conv3x3(cin, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        return self.relu(self.bn2(self.conv2(out)))


class CifarNet(nn.Module):
    """6n+2 layers. `residual=False` gives the plain counterpart."""

    def __init__(self, n=3, residual=True, num_classes=10):
        super().__init__()
        block = BasicBlock if residual else PlainBlock
        self.residual = residual

        self.conv1 = conv3x3(3, 16)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)

        self.cin = 16
        self.stage1 = self._stage(block, 16, n, 1)
        self.stage2 = self._stage(block, 32, n, 2)
        self.stage3 = self._stage(block, 64, n, 2)

        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(64, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _stage(self, block, planes, blocks, stride):
        downsample = None
        if self.residual and (stride != 1 or self.cin != planes):
            # option A: parameter-free, so plain and residual nets match exactly
            downsample = ZeroPadShortcut(stride, planes - self.cin)

        layers = [block(self.cin, planes, stride, downsample)]
        self.cin = planes
        layers += [block(planes, planes) for _ in range(1, blocks)]
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.stage3(self.stage2(self.stage1(x)))
        return self.fc(torch.flatten(self.pool(x), 1))


def depth_to_n(layers):
    assert (layers - 2) % 6 == 0, "CIFAR ResNet depth must be 6n+2"
    return (layers - 2) // 6

The experiment runner

The paper's exact recipe: SGD with momentum 0.9, weight decay 1e-4, batch size 128, learning rate 0.1 divided by 10 at 32k and 48k iterations, stopping at 64k. Augmentation is 4-pixel padding, a random 32x32 crop, and a horizontal flip.

pythonexperiments/degradation.py
import json
import time

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from resnet.cifar import CifarNet, depth_to_n

MEAN, STD = (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)
BATCH, TOTAL_ITERS, DROPS = 128, 64_000, (32_000, 48_000)


def device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def loaders(root="./data"):
    train_tf = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),
    ])
    test_tf = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),
    ])
    tr = datasets.CIFAR10(root, train=True, download=True, transform=train_tf)
    te = datasets.CIFAR10(root, train=False, download=True, transform=test_tf)
    # num_workers=0 on Apple Silicon; MPS does not like worker processes
    nw = 0 if device().type == "mps" else 4
    return (
        DataLoader(tr, BATCH, shuffle=True, num_workers=nw, drop_last=True),
        DataLoader(te, 256, shuffle=False, num_workers=nw),
    )


@torch.no_grad()
def error_rate(model, loader, dev):
    model.eval()
    wrong = total = 0
    for x, y in loader:
        pred = model(x.to(dev)).argmax(1).cpu()
        wrong += (pred != y).sum().item()
        total += y.numel()
    return 100.0 * wrong / total


def train_one(layers, residual, dev, train_loader, test_loader, log_every=2_000):
    model = CifarNet(depth_to_n(layers), residual=residual).to(dev)
    params = sum(p.numel() for p in model.parameters())
    name = f"{'ResNet' if residual else 'plain'}-{layers}"
    print(f"\n{name}: {params/1e6:.3f}M parameters")

    opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=list(DROPS), gamma=0.1)
    criterion = nn.CrossEntropyLoss()

    history, it, start = [], 0, time.time()
    while it < TOTAL_ITERS:
        model.train()
        for x, y in train_loader:
            if it >= TOTAL_ITERS:
                break
            x, y = x.to(dev), y.to(dev)
            opt.zero_grad(set_to_none=True)
            out = model(x)
            loss = criterion(out, y)
            loss.backward()
            opt.step()
            sched.step()

            # running training error on this batch
            train_err = 100.0 * (out.argmax(1) != y).float().mean().item()
            it += 1

            if it % log_every == 0:
                te = error_rate(model, test_loader, dev)
                history.append({"iter": it, "train_err": train_err, "test_err": te})
                print(f"  iter {it:>6}  train {train_err:5.2f}%  test {te:5.2f}%")
                model.train()

    final = error_rate(model, test_loader, dev)
    print(f"  final test error {final:.2f}%  ({time.time()-start:.0f}s)")
    return {"name": name, "layers": layers, "residual": residual,
            "params": params, "final_test_err": final, "history": history}


if __name__ == "__main__":
    dev = device()
    print(f"device: {dev}")
    train_loader, test_loader = loaders()

    results = []
    # The comparison that matters. Plain-56 should end up WORSE than plain-20,
    # on training error as well as test error. ResNet-56 should beat ResNet-20.
    for layers, residual in [(20, False), (56, False), (20, True), (56, True)]:
        results.append(train_one(layers, residual, dev, train_loader, test_loader))

    with open("degradation_results.json", "w") as f:
        json.dump(results, f, indent=2)

    print("\n" + "=" * 52)
    print(f"{'model':<14}{'params':>10}{'test error':>14}")
    for r in results:
        print(f"{r['name']:<14}{r['params']/1e6:>9.3f}M{r['final_test_err']:>13.2f}%")
    print("\nPaper reference: ResNet-20 8.75%, ResNet-56 6.97%.")
    print("Plain nets are not tabulated in the paper; Figure 6 shows them")
    print("rising with depth, roughly 9% at 20 layers to 11-12% at 56.")

What you are looking for is not the absolute numbers, which depend on your seed and setup. It is the sign of the slope. Going from 20 to 56 layers should help the residual network and hurt the plain one, using the same data, the same recipe, and the same parameter count.

On an M4 Max expect roughly 20 to 40 minutes per model, so under three hours for all four. On a single modern NVIDIA GPU expect about 10 minutes each. See What it costs now for where those numbers come from.

Checking the paper's claim about gradients

The paper asserts that degradation is not caused by vanishing gradients, because with batch norm "the backward propagated gradients exhibit healthy norms." That is a checkable claim, and checking it is what separates understanding the paper from reciting it.

pythonexperiments/gradient_norms.py
import torch
import torch.nn as nn

from resnet.cifar import CifarNet, depth_to_n


def per_stage_grad_norms(model, x, y):
    model.train()
    model.zero_grad(set_to_none=True)
    nn.CrossEntropyLoss()(model(x), y).backward()

    norms = {}
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d) and module.weight.grad is not None:
            norms[name] = module.weight.grad.norm().item()
    return norms


if __name__ == "__main__":
    torch.manual_seed(0)
    x, y = torch.randn(128, 3, 32, 32), torch.randint(0, 10, (128,))

    for residual in (False, True):
        model = CifarNet(depth_to_n(56), residual=residual)
        norms = per_stage_grad_norms(model, x, y)
        vals = list(norms.values())
        label = "ResNet-56" if residual else "plain-56"
        print(f"\n{label}: {len(vals)} conv layers")
        print(f"  first layer grad norm: {vals[0]:.4e}")
        print(f"  last  layer grad norm: {vals[-1]:.4e}")
        print(f"  ratio last/first:      {vals[-1]/vals[0]:.2f}")
        print(f"  min {min(vals):.3e}  max {max(vals):.3e}")

Run it and you will find what the paper reported: with batch norm in place, gradients reach the first layer of a 56-layer plain network at a perfectly workable magnitude. Nothing has vanished. The network still trains to a worse solution than its 20-layer counterpart. That is why the paper had to introduce a new name, degradation, rather than reusing the vanishing gradient story.

Training the real thing on ImageNet

For completeness, the paper's ImageNet recipe. See What it costs now before starting this; it is about 13 hours on one A100 and several days on a Mac.

pythonexperiments/imagenet.py
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from resnet.model import resnet50

BATCH, EPOCHS, BASE_LR = 256, 90, 0.1


def build_loaders(root, workers=16):
    normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    train_tf = transforms.Compose([
        # scale augmentation: shorter side sampled from [256, 480], per the paper
        transforms.RandomResizedCrop(224, scale=(0.08, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(0.4, 0.4, 0.4),
        transforms.ToTensor(),
        normalize,
    ])
    val_tf = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        normalize,
    ])
    tr = datasets.ImageFolder(f"{root}/train", train_tf)
    va = datasets.ImageFolder(f"{root}/val", val_tf)
    return (
        DataLoader(tr, BATCH, shuffle=True, num_workers=workers, pin_memory=True,
                   drop_last=True, persistent_workers=True),
        DataLoader(va, BATCH, shuffle=False, num_workers=workers, pin_memory=True),
    )


def main(root="/data/imagenet"):
    dev = torch.device("cuda")
    model = resnet50().to(dev).to(memory_format=torch.channels_last)

    # no weight decay on BatchNorm parameters: a standard later refinement
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        (no_decay if p.ndim <= 1 else decay).append(p)
    opt = torch.optim.SGD(
        [{"params": decay, "weight_decay": 1e-4},
         {"params": no_decay, "weight_decay": 0.0}],
        lr=BASE_LR, momentum=0.9, nesterov=True,
    )

    train_loader, val_loader = build_loaders(root)
    steps = EPOCHS * len(train_loader)
    sched = torch.optim.lr_scheduler.OneCycleLR(
        opt, max_lr=BASE_LR, total_steps=steps, pct_start=0.03, anneal_strategy="cos"
    )
    scaler = torch.amp.GradScaler("cuda")
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

    for epoch in range(EPOCHS):
        model.train()
        for x, y in train_loader:
            x = x.to(dev, non_blocking=True, memory_format=torch.channels_last)
            y = y.to(dev, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = criterion(model(x), y)
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
            sched.step()
        torch.save({"epoch": epoch, "model": model.state_dict()}, "checkpoint.pt")


if __name__ == "__main__":
    main()

This deviates from 2015 in three places, all of which are now standard and all of which help: mixed precision, a cosine schedule with warmup instead of hand-placed step drops, and label smoothing. The original recipe used plain step decay at plateaus, fp32, and no smoothing.

Running on Apple Silicon

pythontools/mps_notes.py
import torch

# bf16 autocast works on MPS in current PyTorch and is the single biggest win
device = torch.device("mps")
model = model.to(device)

for x, y in loader:
    x, y = x.to(device), y.to(device)
    with torch.autocast("mps", dtype=torch.bfloat16):
        loss = criterion(model(x), y)
    loss.backward()
    opt.step()
    opt.zero_grad(set_to_none=True)

Practical notes, learned the hard way by everyone who has tried this:

Using the pretrained model instead

Worth saying plainly: unless you are studying the optimization behaviour, download the weights. Eleven years of downstream work assumes you did.

python
import torch
from torchvision.models import resnet50, ResNet50_Weights

weights = ResNet50_Weights.IMAGENET1K_V2   # 80.86% top-1, improved recipe
model = resnet50(weights=weights).eval()

preprocess = weights.transforms()
batch = preprocess(image).unsqueeze(0)
probs = model(batch).softmax(dim=1)
top5 = probs.topk(5)
for score, idx in zip(top5.values[0], top5.indices[0]):
    print(f"{weights.meta['categories'][idx]:<28}{score.item():.3f}")

Note IMAGENET1K_V2 at 80.86% top-1 against the original recipe's roughly 76%. The architecture did not change at all. The gap is entirely better training: longer schedules, cosine learning rates, label smoothing, mixup, and improved augmentation. That is a useful reminder that the 2015 numbers measure a 2015 recipe, not a ceiling on the architecture.