Domain adaptation for vision

Setup#

pip install torch numpy

Written and run against torch 2.13.0 (CPU) and numpy 2.2.6. No downloads. The whole script trains a small CNN from scratch on synthetic shapes and runs in a few seconds.

A camera change, and a three-line repair#

import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader

rng = np.random.default_rng(0); torch.manual_seed(0)

def draw(cls, bright, contrast):
    img = rng.normal(0.35, 0.06, (32, 32))
    r, c = rng.integers(9, 23, 2); s = int(rng.integers(4, 7))
    yy, xx = np.ogrid[:32, :32]; d = np.hypot(yy - r, xx - c)
    m = [(abs(yy-r) < s) & (abs(xx-c) < s), d < s, abs(d-s) < 1.2,
         (abs(yy-r) < 1.2) | (abs(xx-c) < 1.2)][cls]
    img[m] = 0.95
    return np.clip((img - 0.5) * contrast + 0.5 + bright, 0, 1).astype("float32")

def make(n, bright=0.0, contrast=1.0):
    lab = rng.integers(0, 4, n)
    x = torch.tensor(np.stack([draw(c, bright, contrast) for c in lab])).unsqueeze(1)
    return x, torch.tensor(lab)

# SOURCE = the camera you trained on.  TARGET = a dimmer, flatter camera.
Xs, ys = make(1200)
Xs_te, ys_te = make(400)
Xt, yt = make(400, bright=-0.18, contrast=0.45)
Xt_pool, _ = make(400, bright=-0.18, contrast=0.45)     # unlabelled target images

net = nn.Sequential(
    nn.Conv2d(1,16,3,padding=1), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(16,32,3,padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
    nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32,4))
opt = torch.optim.AdamW(net.parameters(), lr=3e-3)
dl = DataLoader(TensorDataset(Xs, ys), batch_size=64, shuffle=True)
for _ in range(25):
    net.train()
    for xb, yb in dl:
        opt.zero_grad(); F.cross_entropy(net(xb), yb).backward(); opt.step()

def acc(x, y):
    net.eval()
    with torch.no_grad():
        return (net(x).argmax(1) == y).float().mean().item()

print(f"same camera as training  : {acc(Xs_te, ys_te):.3f}")
print(f"different camera         : {acc(Xt, yt):.3f}")

# AdaBN: recompute BatchNorm statistics on UNLABELLED target images. No labels, no gradients.
for mod in net.modules():
    if isinstance(mod, nn.BatchNorm2d):
        mod.reset_running_stats(); mod.momentum = None    # None = plain running average
net.train()
with torch.no_grad():
    for i in range(0, 400, 64):
        net(Xt_pool[i:i+64])
print(f"different camera + AdaBN : {acc(Xt, yt):.3f}")

same camera as training  : 1.000
different camera         : 0.275
different camera + AdaBN : 1.000

Deterministic across runs with the seeds set. Your numbers should match. If the middle line is not far below the first, check the brightness and contrast shift.

What that run shows#

1.000 to 0.275 is a collapse, and the shift was mild. The target images are dimmer and flatter — brightness down 0.18, contrast scaled to 0.45. A human sees the same four shapes without effort. Four-class chance is 0.250, so the model finished barely above guessing.

Zero labels were used to fix it. Xt_pool was created, its labels discarded, and only the images passed through. This matters commercially: you always have unlabelled production images, and you rarely have labelled ones.

No weight moved. torch.no_grad() was on, no optimiser ran. Only the BatchNorm buffers changed. The classifier that scored 0.275 and the classifier that scored 1.000 have identical weights.

That is the point this lesson exists to make. The knowledge was never lost. It was being read through a wrong scale.

Line by line#

mod.reset_running_stats() clears running_mean and running_var and sets num_batches_tracked to zero.

mod.momentum = None switches BatchNorm from an exponential moving average to a cumulative average over every batch it sees. With the default momentum of 0.1, the final statistics would be dominated by the last few batches. Setting it to None uses all 400 target images equally.

net.train() is required and looks wrong. Training mode is what makes BatchNorm update its buffers. See freezing and unfreezing layers for the same mechanism appearing as a bug. Here it is the feature.

acc() calls net.eval() internally, so measurement uses the newly stored statistics rather than the test batch's own.

When AdaBN will not save you#

There is also a hard requirement: your architecture must contain BatchNorm. Vision transformers use LayerNorm. Its statistics are computed per sample at inference, so they need no adaptation and get no free repair.

Common mistakes#

Adapting on a handful of target images. BatchNorm statistics estimated from 30 images are noisy, and you can end up worse than before. Fix: use several hundred, and check on any labels you do have.

Adapting on a target batch that is one class. If the production stream is sorted, a batch of 64 identical objects gives statistics for that object, not the domain. Fix: shuffle across a large buffer before adapting.

Adapting once and forgetting. Domains keep moving — seasons, lens dirt, a replaced bulb. Fix: recalibrate on a schedule and log the statistics so you can see drift. See monitoring and drift.

Calling it domain adaptation when the labels changed. If the new site also defines "defective" differently, no unsupervised method fixes that. Fix: label a small target set and measure, before assuming the problem is visual.

Try it yourself#

Change contrast=0.45 to contrast=1.0 and keep only the brightness shift. Measure how much of the drop AdaBN recovers. Then shift the shapes' sizes instead of their brightness, and watch AdaBN stop helping. That is a content change wearing a domain-shift costume.

What to learn next#

• Geometric augmentations — building tolerance to shift into training, before you need to adapt.
• BatchNorm in PyTorch — the buffers this whole lesson turns on.
• Monitoring and drift — noticing a domain shift before your users do.