PyTorch Mixup
Mixup is training on blends of two inputs, scored against blends of their two labels. The blends teach the model softer boundaries between classes.
A blend of two images
Mixup adds two images with weights lam and 1 - lam. A lam of 0.5 gives a half-and-half ghost of both shapes.
blend = 0.5 * train_x[0] + 0.5 * train_x[1]
print(blend.shape)
print(f"Largest blend pixel: {float(blend.max()):.1f}")The loss blends with the same lam
The blended image has no single true label. The loss is the same blend of the losses against both labels.
lam = 0.25
xb, yb = train_x[:8], train_y[:8]
perm = torch.tensor([1, 0, 3, 2, 5, 4, 7, 6])
mixed = lam * xb + (1 - lam) * xb[perm]
model = plain
mixed_loss = lam * loss_fn(model(mixed), yb) + (1 - lam) * loss_fn(model(mixed), yb[perm])
print(f"Blended loss: {float(mixed_loss):.4f}")Lam runs from 0 to 1
A lam near 1 keeps the first image almost whole. Blending the first image with a blank image shows the mean pixel sliding with lam.
blank = torch.zeros_like(train_x[0])
for lam in (0.0, 0.5, 1.0):
mixed = lam * train_x[0] + (1 - lam) * blank
print(f"lam {lam}: mean pixel {float(mixed.mean()):.3f}")
print(f"First image mean: {float(train_x[0].mean()):.3f}")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)
rotated_x = torch.rot90(test_x, 1, dims=(2, 3))
diagonal_x, diagonal_y = make_diagonal(90, seed=99)
shifted_x = torch.cat([rotated_x, diagonal_x])
shifted_y = torch.cat([test_y, diagonal_y])
loss_fn = nn.CrossEntropyLoss()
def accuracy(model, x, y):
model.eval()
with torch.no_grad():
predictions = model(x).argmax(dim=1)
return float((predictions == y).float().mean())
def train_model(use_mixup):
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)
mixup_gen = torch.Generator().manual_seed(5)
for epoch in range(12):
model.train()
for xb, yb in loader:
optimizer.zero_grad()
if use_mixup:
lam = float(torch.rand((), generator=mixup_gen))
perm = torch.randperm(len(yb), generator=mixup_gen)
mixed = lam * xb + (1 - lam) * xb[perm]
loss = (lam * loss_fn(model(mixed), yb)
+ (1 - lam) * loss_fn(model(mixed), yb[perm]))
else:
loss = loss_fn(model(xb), yb)
loss.backward()
optimizer.step()
return model
plain = train_model(False)
mixed_model = train_model(True)
print(f"Plain run: clean accuracy {accuracy(plain, test_x, test_y):.3f}, "
f"rotated accuracy {accuracy(plain, shifted_x, shifted_y):.3f}")
print(f"Mixup run: clean accuracy {accuracy(mixed_model, test_x, test_y):.3f}, "
f"rotated accuracy {accuracy(mixed_model, shifted_x, shifted_y):.3f}")Key points
- Mixup: each training batch is a weighted blend of two batches, and its loss blends the same way
- Lam: the blend weight, drawn fresh for each batch between 0 and 1
- A lam draw from the uniform range matches mixup's alpha setting of 1.0
- The shuffled partner batch comes from a seeded randperm of the same batch
- Blended inputs hold values between 0 and 1 wherever the two shapes disagree
- Both runs in the closing example share data, seeds, epochs, and batch order
In this example, two shapes CNNs train on the same upright data for 12 epochs. The plain run trains on the batches as they come. The mixup run blends each batch with a shuffled copy of itself and blends the loss with the same lam. Lam comes from a seeded generator, so both runs stay repeatable. Both models score 1.000 on the clean test set. The rotated test holds 180 images, half turned by rot90 and half redrawn with diagonal bars. The plain model scores 0.728 there. The mixup model scores 0.833. Training on ghosts of two shapes at once left the second model less tied to exact pixel patterns.
Note
Mixup never shows the model a rotated shape here, and the rotated win still appears.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
What does the mixup run's loss blend, and with what weights?