Project: image augmentation and head training¶
Connect image variation, impulse noise and a replacement classification head while keeping the backbone fixed.
What to remember¶
Transfer learning reuses pretrained features; head-only training learns a new output mapping. Freezing parameters alone does not freeze BatchNorm statistics. This script holds the backbone in evaluation mode and trains only the new head.
The default run uses synthetic images and a random ResNet backbone, so it demonstrates the mechanics without downloads. The optional --data … --pretrained run uses ImageNet weights and your own image folders for actual transfer learning.
Important code¶
model = resnet18(weights=weights)
model.requires_grad_(False)
model.fc = nn.Linear(model.fc.in_features, len(classes))
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
# Repeat at the start of each head-only training epoch:
model.eval() # keep backbone BatchNorm statistics fixed
model.fc.train() # the new head is the part being trained
For your image folders, the training transform follows this order:
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
ImpulseNoise(noise),
transforms.Normalize(mean, std),
])
ImpulseNoise(0.02) replaces about 2% of spatial pixels with black or white in expectation, using one mask shared by RGB channels. It runs on floats in 0–1, before normalization. The amount varies by sample. This is the salt-and-pepper mechanism studied in the course; it is not a complete camera-noise model.
| Decision | Training | Validation |
|---|---|---|
| Geometry | Random crop and optional flip | Fixed resize and centre crop |
| Noise | Optional selected severity | Clean, deterministic input |
| Normalization | Matching channel mean/std | The same channel mean/std |
| Layers | New fc head learns |
All layers evaluate |
Flips/crops must preserve your labels. For a separate robustness evaluation, create a fixed noisy validation set rather than randomizing ordinary validation.
Visual noise check¶
Run and inspect¶
This small CPU run writes report.json and head-model.pt under artifacts/vision-head/. The report explicitly records random weights and synthetic data. Its accuracy has no transfer-quality meaning.
For real transfer learning, provide separate my-images/train/class-name/ and my-images/val/class-name/ folders containing the same class names:
The first pretrained run may download weights. The example stays on CPU and zero DataLoader workers for portability. The real-image/pretrained branch is provided for study; the local verification used offline data, not a new transfer-learning benchmark.
The checkpoint saves the head and backbone state together with class names, validation preprocessing, original backbone weight identity and noise policy. After adapting a head, ImageNet's original class names no longer describe its outputs.
Try changing: compare --noise 0 with --noise 0.02 on the same split. Do not expect noise to help every task. To fine-tune late layers, explicitly unfreeze them and rebuild the optimizer.
Review transforms and noise · Review transfer stages
What the retained execution shows¶
Seeded transform execution on an original illustration. The panels apply actual TorchVision crop and colour transforms, followed by the maintained ImpulseNoise code. This scene is for inspecting transforms; the head-training run uses seeded random images. Training-only transformations preserve labels only when the selected change is plausible for your task. Exact transform order, shape and ranges.
The retained two-epoch CPU head run uses random ResNet18 weights, 24 synthetic training images, 12 validation images and three synthetic labels. Random pixels have no useful class signal; its validation score does not demonstrate transfer learning.
| Before/after check | Observed in saved state |
|---|---|
| Frozen backbone parameters (11,176,512 values) | Unchanged |
| BatchNorm running buffers | Unchanged: backbone remains in evaluation mode |
| Trainable classifier (1,539 values) | Changed after optimizer updates |
The script verifies these checks against a cloned initial state. Retained head report. To learn useful features for your own classes, supply labelled train/ and val/ folders; --pretrained optionally downloads ImageNet weights and uses their preprocessing contract. A random frozen backbone is a mechanics demo, not pretrained transfer learning.
Complete source¶
Open the maintained runnable script
"""Head-only ResNet training with noise; offline demo or your ImageFolder data."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
from torchvision.models import ResNet18_Weights, resnet18
from recall_patterns import ImpulseNoise
class ToyImages(Dataset):
"""Original synthetic inputs to check mechanics, not transfer quality."""
def __init__(self, size, transform, seed):
generator = torch.Generator().manual_seed(seed)
self.images = torch.rand(size, 3, 32, 32, generator=generator)
self.labels = torch.arange(size) % 3
self.transform = transform
def __len__(self):
return len(self.labels)
def __getitem__(self, index):
return self.transform(self.images[index]), self.labels[index]
def run(output: Path, data=None, pretrained=False, epochs=2, noise=0.02):
torch.set_num_threads(2)
torch.manual_seed(31)
weights = ResNet18_Weights.DEFAULT if pretrained else None
if pretrained and data is None:
raise ValueError("--pretrained needs --data with train/ and val/ class folders")
if data is None:
train_data = ToyImages(24, transforms.Compose([
transforms.RandomHorizontalFlip(), ImpulseNoise(noise)]), 41)
val_data = ToyImages(12, transforms.Lambda(lambda image: image.clone()), 42)
classes = ["synthetic-0", "synthetic-1", "synthetic-2"]
preprocessing = {"size": 32, "range": [0, 1], "normalization": None}
else:
# PIL -> augmentation -> float [0,1] -> noise -> normalization.
if weights is not None:
mean, std = weights.transforms().mean, weights.transforms().std
val_transform = weights.transforms()
else:
mean, std = [0.5] * 3, [0.5] * 3
val_transform = transforms.Compose([
transforms.Resize(256), transforms.CenterCrop(224),
transforms.ToTensor(), transforms.Normalize(mean, std)])
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(),
transforms.ToTensor(), ImpulseNoise(noise), transforms.Normalize(mean, std)])
train_data = datasets.ImageFolder(data / "train", transform=train_transform)
val_data = datasets.ImageFolder(data / "val", transform=val_transform)
if train_data.class_to_idx != val_data.class_to_idx:
raise ValueError("train/ and val/ must contain the same class folders")
classes = train_data.classes
preprocessing = {"size": 224, "validation": "resize 256 + center crop 224",
"mean": mean, "std": std, "range_before_normalization": [0, 1]}
train_loader = DataLoader(train_data, batch_size=8, shuffle=True)
val_loader = DataLoader(val_data, batch_size=8)
model = resnet18(weights=weights)
model.requires_grad_(False)
model.fc = nn.Linear(model.fc.in_features, len(classes))
before = {name: value.detach().clone() for name, value in model.state_dict().items()}
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
history = []
for epoch in range(epochs):
# Frozen backbone includes BatchNorm buffers, not only parameter gradients.
model.eval()
model.fc.train()
for images, labels in train_loader:
optimizer.zero_grad(set_to_none=True)
loss = nn.functional.cross_entropy(model(images), labels)
loss.backward()
optimizer.step()
model.eval()
correct = 0
with torch.no_grad():
for images, labels in val_loader:
correct += (model(images).argmax(1) == labels).sum().item()
history.append({"epoch": epoch + 1, "validation_accuracy": correct / len(val_data)})
output.mkdir(parents=True, exist_ok=True)
report = {
"mode": "own-image-data" if data is not None else "synthetic mechanism demo",
"backbone_weights": weights.name if weights is not None else "random",
"classes": classes, "preprocessing": preprocessing,
"noise_amount": noise, "train_samples": len(train_data),
"validation_samples": len(val_data),
"trainable_parameters": sum(p.numel() for p in model.parameters() if p.requires_grad),
"frozen_parameter_count": sum(p.numel() for p in model.parameters() if not p.requires_grad),
"frozen_parameters_unchanged": all(torch.equal(before[name], p) for name, p in model.named_parameters() if not p.requires_grad),
"frozen_buffers_unchanged": all(torch.equal(before[name], b) for name, b in model.named_buffers()),
"head_changed": any(not torch.equal(before[name], p) for name, p in model.named_parameters() if name.startswith("fc.")),
"history": history,
}
torch.save({"state_dict": model.state_dict(), "metadata": report}, output / "head-model.pt")
(output / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
print(f"{report['mode']} / {report['backbone_weights']} backbone; "
f"{report['trainable_parameters']} trainable parameters")
return report
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--output", type=Path, default=Path("artifacts/vision-head"))
parser.add_argument("--data", type=Path, help="Root containing train/class/ and val/class/")
parser.add_argument("--pretrained", action="store_true", help="May download ImageNet weights")
parser.add_argument("--epochs", type=int, default=2)
parser.add_argument("--noise", type=float, default=0.02)
args = parser.parse_args()
if args.epochs < 1 or not 0 <= args.noise <= 1:
parser.error("Use positive --epochs and --noise in [0,1]")
if args.pretrained and args.data is None:
parser.error("--pretrained needs --data; offline demo intentionally uses random weights")
run(args.output, args.data, args.pretrained, args.epochs, args.noise)
Open the shared impulse-noise transform
class ImpulseNoise:
"""Replace selected spatial pixels in a float [C,H,W] image in [0,1]."""
def __init__(self, amount=0.02, salt_fraction=0.5):
if not 0 <= amount <= 1 or not 0 <= salt_fraction <= 1:
raise ValueError("Noise probabilities must be in [0, 1]")
self.amount = amount
self.salt_fraction = salt_fraction
def __call__(self, image):
# One mask shared across channels: a selected RGB pixel is white or black.
draw = torch.rand_like(image[:1])
salt = draw < self.amount * self.salt_fraction
pepper = (draw >= self.amount * self.salt_fraction) & (draw < self.amount)
noisy = torch.where(salt, torch.ones_like(image), image)
return torch.where(pepper, torch.zeros_like(image), noisy)
API reference: ResNet18 and its weight preprocessing.
Visual-generation source · runnable with the retained reports
"""Render original recall figures from retained reports; never replace legacy results."""
from pathlib import Path
import json
import sys
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from torchvision import transforms
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "examples"))
from recall_patterns import ImpulseNoise
def render():
plt.rcParams.update({"font.size": 12})
reports = ROOT / "docs/assets/data/recall-2026-10-01"
images = ROOT / "docs/assets/images"
training = json.loads((reports / "comparison.json").read_text())
for key, title, ylabel, filename in [
("loss", "Validation loss · 48 synthetic samples", "Cross-entropy", "training-loss.png"),
("macro_f1", "Validation macro F1 · 3 classes", "Macro F1", "training-f1.png")]:
fig, ax = plt.subplots(figsize=(4.5, 3.6))
for trial in training["trials"]:
ax.plot([h["epoch"] for h in trial["history"]], [h[key] for h in trial["history"]], marker="o", label=f"Initial LR {trial['initial_lr']}")
ax.set(title=title, xlabel="Epoch", ylabel=ylabel)
ax.legend(); ax.grid(alpha=.2); fig.tight_layout()
fig.savefig(images / filename, dpi=160); plt.close(fig)
best = next(t for t in training["trials"] if t["initial_lr"] == training["best_initial_lr"])
matrix = np.array(best["history"][-1]["confusion_matrix"])
fig, ax = plt.subplots(figsize=(5, 4))
chart = ax.imshow(matrix, cmap="Blues")
for (row, col), value in np.ndenumerate(matrix):
ax.text(col, row, str(value), ha="center", va="center", color="white" if value > matrix.max()/2 else "black", fontsize=15)
ax.set(title=f"Final validation · initial LR {best['initial_lr']}", xlabel="Predicted class", ylabel="True class", xticks=range(3), yticks=range(3))
fig.colorbar(chart, ax=ax, label="Samples"); fig.tight_layout()
fig.savefig(images / "training-confusion.png", dpi=160); plt.close(fig)
losses = [1, .8, .8, .8, .7, .7, .7, .7, .7, .7, .7, .7]
fig, ax = plt.subplots(figsize=(6, 4))
schedules = {}
for name in ["StepLR", "CosineAnnealingLR", "ReduceLROnPlateau"]:
parameter = torch.nn.Parameter(torch.zeros(1))
optimizer = torch.optim.SGD([parameter], lr=.1)
scheduler = (torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=.5) if name == "StepLR" else
torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=12) if name == "CosineAnnealingLR" else
torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=.5, patience=1))
rates = []
for loss in losses:
rates.append(optimizer.param_groups[0]["lr"])
optimizer.step()
scheduler.step(loss) if name == "ReduceLROnPlateau" else scheduler.step()
schedules[name] = rates
ax.plot(range(1, 13), rates, marker="o", label=name)
ax.set(title="Toy schedule calls · initial LR 0.1", xlabel="Epoch (LR used before scheduler call)", ylabel="Learning rate")
ax.legend(fontsize=10); ax.grid(alpha=.2); fig.tight_layout()
svg = images / "training-schedules.svg"
fig.savefig(svg); plt.close(fig)
svg.write_text("\n".join(line.rstrip() for line in svg.read_text(encoding="utf-8").splitlines()) + "\n", encoding="utf-8")
(reports / "schedules.json").write_text(json.dumps({"kind": "toy scheduler execution, not training", "validation_loss_signal": losses, "rates_used": schedules}, indent=2), encoding="utf-8")
# An original inspection scene, not a photograph or a course image.
y, x = np.mgrid[:180, :220]
pixels = np.stack([70 + x*.3, 100 + y*.3, 150 + x*.1], axis=-1).astype(np.uint8)
pixels[25:155, 25:195] = [175, 191, 205]
pixels[65:115, 145:150] = [55, 70, 85]
pixels[(x-75)**2 + (y-75)**2 < 140] = [75, 90, 100]
original = Image.fromarray(pixels)
torch.manual_seed(71)
crop = transforms.RandomResizedCrop((180, 220), scale=(.65, .85))(original)
colour = transforms.ColorJitter(brightness=.3, contrast=.3, saturation=.3)(crop)
noisy = ImpulseNoise(.05)(transforms.ToTensor()(colour))
panels = [np.asarray(original), np.asarray(crop), np.asarray(colour), noisy.permute(1, 2, 0).numpy()]
# 2x2 layout stays legible when scaled to a mobile column.
fig, axes = plt.subplots(2, 2, figsize=(6, 5.5))
for ax, panel, title in zip(axes.flat, panels, ["Original illustration", "RandomResizedCrop", "ColorJitter after crop", "ImpulseNoise (p = 0.05)"]):
ax.imshow(panel); ax.set_title(title, fontsize=15); ax.axis("off")
fig.tight_layout(); fig.savefig(images / "vision-transforms.png", dpi=160); plt.close(fig)
normalized = transforms.Normalize([.5]*3, [.5]*3)(noisy)
report = {"kind": "seeded execution on an original illustration", "seed": 71,
"transforms": ["RandomResizedCrop: scale .65–.85, output 180×220", "ColorJitter: brightness/contrast/saturation .3", "ToTensor", "ImpulseNoise .05", "Normalize mean=.5 std=.5"],
"float_shape": list(noisy.shape), "float_range": [noisy.min().item(), noisy.max().item()],
"normalized_range": [normalized.min().item(), normalized.max().item()]}
(reports / "transforms.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print("Rendered training loss, F1, confusion matrix and seeded transform panels from retained evidence")
if __name__ == "__main__":
render()
Run python scripts/render_recall_evidence.py from the repository root to regenerate these new explanatory charts and transform panels. It reads the public retained reports and leaves the legacy regression and EMNIST images intact.