PyTorch When the World Shifts

A distribution shift is a change between the data a model trained on and the data it meets later. Rotating the test shapes is one shift, and the accuracy drop measures it.

rot90 turns the whole image

torch.rot90 rotates a tensor by a quarter turn. The dims argument names the two axes of the image plane.

tiny = torch.zeros(3, 3)
tiny[0, :] = 1.0
print(torch.rot90(tiny, 1, dims=(0, 1)))

A diagonal bar from a distance rule

No tensor op rotates pixels by 45 degrees. The dataset draws the bar again along the diagonal instead.

bar = draw_diagonal_bar(torch.Generator().manual_seed(3))
print(f"Diagonal bar cells: {int(bar.sum())}")
print(f"Horizontal bar cells: {int(draw_shape(0, torch.Generator().manual_seed(3)).sum())}")

Two shapes that rotation cannot hurt

A ring looks the same after a quarter turn, and so does a square. Their filled-cell counts cannot change under rot90.

ring = draw_shape(1, torch.Generator().manual_seed(4))
turned = torch.rot90(ring, 1, dims=(0, 1))
print(f"Ring cells before: {int(ring.sum())}, after: {int(turned.sum())}")

Example

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def draw_shape(label, g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    if label == 0:
        centre = 8 + int(torch.randint(-2, 3, (1,), generator=g))
        mask = (rows >= centre - 1) & (rows <= centre + 1) & (cols >= 0)
    elif label == 1:
        cr = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        cc = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        dist = torch.sqrt((rows - cr) ** 2 + (cols - cc) ** 2)
        mask = (dist >= 4.5) & (dist <= 6.5)
    else:
        top = int(torch.randint(2, 7, (1,), generator=g))
        left = int(torch.randint(2, 7, (1,), generator=g))
        mask = (rows >= top) & (rows <= top + 7) & (cols >= left) & (cols <= left + 7)
    return mask.float()


def make_shapes(count, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.arange(count) % 3
    labels = labels[torch.randperm(count, generator=g)]
    images = [draw_shape(int(label), g) for label in labels]
    return torch.stack(images).unsqueeze(1), labels


class ShapesCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Linear(16, 3)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = torch.relu(self.conv2(x))
        x = self.gap(x)
        return self.head(torch.flatten(x, 1))


def draw_diagonal_bar(g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    shift = int(torch.randint(-3, 4, (1,), generator=g))
    return (((rows - cols) - shift).abs() <= 1.2).float()


def make_diagonal(count, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.arange(count) % 3
    labels = labels[torch.randperm(count, generator=g)]
    images = []
    for label in labels:
        if int(label) == 0:
            images.append(draw_diagonal_bar(g))
        else:
            images.append(draw_shape(int(label), g))
    return torch.stack(images).unsqueeze(1), labels


train_x, train_y = make_shapes(240, seed=42)
test_x, test_y = make_shapes(90, seed=99)
loader = DataLoader(TensorDataset(train_x, train_y), batch_size=32,
                    shuffle=True, generator=torch.Generator().manual_seed(0))

torch.manual_seed(1)
model = ShapesCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(8):
    model.train()
    for xb, yb in loader:
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()


def accuracy(x, y):
    model.eval()
    with torch.no_grad():
        predictions = model(x).argmax(dim=1)
    return float((predictions == y).float().mean())


rotated_x = torch.rot90(test_x, 1, dims=(2, 3))
diagonal_x, diagonal_y = make_diagonal(90, seed=99)

print(f"Accuracy at 0 degrees:  {accuracy(test_x, test_y):.3f}")
print(f"Accuracy at 90 degrees: {accuracy(rotated_x, test_y):.3f}")
print(f"Accuracy at 45 degrees: {accuracy(diagonal_x, diagonal_y):.3f}")

model.eval()
with torch.no_grad():
    rotated_predictions = model(rotated_x).argmax(dim=1)
names = ["bar", "ring", "square"]
for label in range(3):
    mask = test_y == label
    score = float((rotated_predictions[mask] == test_y[mask]).float().mean())
    print(f"Rotated {names[label]} accuracy: {score:.3f}")

Key points

Distribution shift: the test data stops matching the training data in some steady way
torch.rot90 turns a tensor by a quarter turn inside the plane named by dims
No plain tensor op rotates pixels by 45 degrees, so the diagonal bars are drawn fresh
A ring and a square look the same after a quarter turn, so rotation cannot hurt them
A bar carries its class in its direction, so rotation strikes the bar class alone
The model trains on upright shapes only, and every rotated set is scored, never trained on

In this example, the shapes CNN trains for 8 epochs on upright shapes only. The first test is the same upright set, and accuracy is 1.000. torch.rot90 then turns every test image by a quarter turn for the 90 degree score. Accuracy falls to 0.667. The 45 degree set is regenerated instead, with its bars drawn along the diagonal. Accuracy there is 0.789. The per-class rows name the cause. Rotated bars score 0.000, because a standing bar matches nothing the model trained on. Rotated rings and squares score 1.000, because a quarter turn leaves them unchanged. The diagonal bars sit between the two cases, and 11 of the 30 are still read as bars.

Note

The shift here is visible and total, while real shifts are often partial and hidden in one feature.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def draw_shape(label, g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    if label == 0:
        centre = 8 + int(torch.randint(-2, 3, (1,), generator=g))
        mask = (rows >= centre - 1) & (rows <= centre + 1) & (cols >= 0)
    elif label == 1:
        cr = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        cc = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        dist = torch.sqrt((rows - cr) ** 2 + (cols - cc) ** 2)
        mask = (dist >= 4.5) & (dist <= 6.5)
    else:
        top = int(torch.randint(2, 7, (1,), generator=g))
        left = int(torch.randint(2, 7, (1,), generator=g))
        mask = (rows >= top) & (rows <= top + 7) & (cols >= left) & (cols <= left + 7)
    return mask.float()


def make_shapes(count, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.arange(count) % 3
    labels = labels[torch.randperm(count, generator=g)]
    images = [draw_shape(int(label), g) for label in labels]
    return torch.stack(images).unsqueeze(1), labels


class ShapesCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Linear(16, 3)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = torch.relu(self.conv2(x))
        x = self.gap(x)
        return self.head(torch.flatten(x, 1))


def draw_diagonal_bar(g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    shift = int(torch.randint(-3, 4, (1,), generator=g))
    return (((rows - cols) - shift).abs() <= 1.2).float()


def make_diagonal(count, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.arange(count) % 3
    labels = labels[torch.randperm(count, generator=g)]
    images = []
    for label in labels:
        if int(label) == 0:
            images.append(draw_diagonal_bar(g))
        else:
            images.append(draw_shape(int(label), g))
    return torch.stack(images).unsqueeze(1), labels


train_x, train_y = make_shapes(240, seed=42)
test_x, test_y = make_shapes(90, seed=99)
loader = DataLoader(TensorDataset(train_x, train_y), batch_size=32,
                    shuffle=True, generator=torch.Generator().manual_seed(0))

torch.manual_seed(1)
model = ShapesCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(8):
    model.train()
    for xb, yb in loader:
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()


def accuracy(x, y):
    model.eval()
    with torch.no_grad():
        predictions = model(x).argmax(dim=1)
    return float((predictions == y).float().mean())


rotated_x = torch.rot90(test_x, 1, dims=(2, 3))
diagonal_x, diagonal_y = make_diagonal(90, seed=99)

print(f"Accuracy at 0 degrees:  {accuracy(test_x, test_y):.3f}")
print(f"Accuracy at 90 degrees: {accuracy(rotated_x, test_y):.3f}")
print(f"Accuracy at 45 degrees: {accuracy(diagonal_x, diagonal_y):.3f}")

model.eval()
with torch.no_grad():
    rotated_predictions = model(rotated_x).argmax(dim=1)
names = ["bar", "ring", "square"]
for label in range(3):
    mask = test_y == label
    score = float((rotated_predictions[mask] == test_y[mask]).float().mean())
    print(f"Rotated {names[label]} accuracy: {score:.3f}")

Exercise

Why do rings and squares keep a perfect score under the 90 degree rotation?

A quarter turn leaves both shapes looking the same The model trained on rotated rings and squares Rotation changes only the labels The test set holds no rings or squares