Train and serve preprocessing mismatch

pip install numpy opencv-python pillow torch

import numpy as np, cv2, torch
import torch.nn.functional as F
from PIL import Image

rng = np.random.default_rng(0)


def make_image(i):
    """A 512x512 RGB picture with fine texture, like real photos have."""
    yy, xx = np.mgrid[0:512, 0:512]
    f = 0.20 + 0.02 * (i % 7)
    base = 0.5 + 0.45 * np.sin(f * xx + 0.3 * i) * np.cos(f * yy * 0.9 + 0.2 * i)
    img = np.stack([base, np.roll(base, 3, 0), np.roll(base, 6, 1)], -1)
    return (np.clip(img, 0, 1) * 255).astype(np.uint8)


def pipe_pil(arr):
    """What most training scripts do: PIL, RGB, bilinear, /255."""
    im = Image.fromarray(arr).resize((224, 224), Image.BILINEAR)
    return np.asarray(im, dtype=np.float32) / 255.0


def pipe_cv2(arr):
    """What most C++/OpenCV serving code does: cv2 bilinear on a BGR array."""
    bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
    small = cv2.resize(bgr, (224, 224), interpolation=cv2.INTER_LINEAR)
    rgb = cv2.cvtColor(small, cv2.COLOR_BGR2RGB)
    return rgb.astype(np.float32) / 255.0


def pipe_torch_nearest(arr):
    """torch.nn.functional.interpolate, whose antialias default is False."""
    t = torch.from_numpy(arr).permute(2, 0, 1).float()[None] / 255.0
    out = F.interpolate(t, size=(224, 224), mode="bilinear",
                        align_corners=False, antialias=False)
    return out[0].permute(1, 2, 0).numpy()


def pipe_torch_aa(arr):
    t = torch.from_numpy(arr).permute(2, 0, 1).float()[None] / 255.0
    out = F.interpolate(t, size=(224, 224), mode="bilinear",
                        align_corners=False, antialias=True)
    return out[0].permute(1, 2, 0).numpy()


PIPES = {"PIL bilinear": pipe_pil,
         "cv2 INTER_LINEAR": pipe_cv2,
         "torch antialias=False": pipe_torch_nearest,
         "torch antialias=True": pipe_torch_aa}

arr = make_image(0)
ref = pipe_pil(arr)
print("difference from the PIL pipeline, on one 512x512 photo resized to 224:")
print(f"{'pipeline':<24} {'mean |diff|':>11} {'max |diff|':>10} {'>1/255':>8} {'cosine':>8}")
for name, fn in PIPES.items():
    out = fn(arr)
    d = np.abs(out - ref)
    cos = float((out.ravel() @ ref.ravel())
                / (np.linalg.norm(out) * np.linalg.norm(ref)))
    print(f"{name:<24} {d.mean():>11.5f} {d.max():>10.5f} "
          f"{(d > 1/255).mean():>7.1%} {cos:>8.5f}")

# A fixed, seeded linear classifier: same weights on every machine and every run.
W = rng.normal(size=(224 * 224 * 3, 5)).astype(np.float32) / 100
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD = np.array([0.229, 0.224, 0.225], np.float32)


def classify(x, normalise=True):
    z = (x - MEAN) / STD if normalise else x
    return int(np.argmax(z.reshape(1, -1) @ W))


imgs = [make_image(i) for i in range(200)]
print("\nout of 200 images, how many change class when the pipeline changes?")
base = [classify(pipe_pil(a)) for a in imgs]
for name, fn in PIPES.items():
    other = [classify(fn(a)) for a in imgs]
    flips = sum(1 for a, b in zip(base, other) if a != b)
    print(f"  trained with PIL, served with {name:<22} {flips:>3}/200")

print("\ntwo mistakes that are not about resizing at all:")
no_norm = sum(1 for a in imgs
              if classify(pipe_pil(a)) != classify(pipe_pil(a), normalise=False))
print(f"  forgot the mean/std normalisation entirely   {no_norm:>3}/200")
swapped = sum(1 for a in imgs
              if classify(pipe_pil(a)) != classify(pipe_pil(a)[..., ::-1].copy()))
print(f"  fed BGR to a model trained on RGB           {swapped:>3}/200")

import torch, torch.nn as nn

class Preprocessed(nn.Module):
    """Takes uint8 NHWC RGB straight from the decoder. No config to forget."""
    def __init__(self, net):
        super().__init__()
        self.net = net
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward(self, x):                       # x: (N, H, W, 3), uint8, RGB
        x = x.permute(0, 3, 1, 2).float() / 255.0
        x = torch.nn.functional.interpolate(
            x, size=(224, 224), mode="bilinear",
            align_corners=False, antialias=True)   # matches transforms.v2.Resize
        return self.net((x - self.mean) / self.std)

net = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(3, 5))
model = Preprocessed(net).eval()

raw = torch.randint(0, 256, (2, 512, 512, 3), dtype=torch.uint8)
with torch.no_grad():
    print("output shape:", tuple(model(raw).shape))
print("buffers travel with the weights:", list(dict(model.named_buffers())))

def test_pipelines_agree():
    img = load_fixture("test_image.png")
    a = training_preprocess(img)
    b = serving_preprocess(img)
    assert np.abs(a - b).max() < 1e-4, f"max diff {np.abs(a - b).max()}"