6 - Few images: class weights, data augmentation and transfer learning¶
In chapter 6 the dermatology project hits the usual problem of medicine: Emma would need several thousand images, and Grace finds only 280 usable ones in the archive, 180 benign, 70 suspicious and 30 malignant, not all of good quality. Emma proposes three remedies: techniques for the class imbalance (oversampling, adjusting the weights), data augmentation (rotations, lighting changes) and transfer learning (starting from a model pretrained on a large dataset and adapting it). Months later, with twice as many images, Noah reports a much better model.
In this notebook we measure each remedy with the book's network (in PyTorch) and 280 synthetic images:
- Import the libraries.
- The images: 280 from Minermont's archive.
- The network and how we measure it.
- The starting point: training from scratch.
- Class imbalance: class weights and oversampling.
- Data augmentation.
- Transfer learning from another hospital's images.
- Twice as many images.
- Summary.
Step 1: Import the libraries¶
We use PyTorch for the network, as in the chapter 7 notebooks; the architecture is the same one the Keras notebook builds.
import copy
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
torch.set_num_threads(8)
Step 2: The images: 280 from Minermont's archive¶
The lesions come from the same generator as the Keras notebook. To imitate an archive of uneven quality, capture gives each photo its own exposure (brightness and contrast), its own amount of noise and, in 40% of the cases, a slight blur. We take 280 training images with the book's split and, to measure the results precisely, 900 test images taken the same way. Step 8 will use 280 more.
def make_lesions(n, seed, shares=(0.5, 0.3, 0.2)):
"""Synthetic 28x28 grayscale images of skin lesions: 0 = benign, 1 = suspicious, 2 = malignant."""
rng = np.random.default_rng(seed)
labels = rng.choice(3, n, p=shares)
yy, xx = np.mgrid[0:28, 0:28].astype(float)
images = np.empty((n, 28, 28), dtype=np.float32)
# the hidden rules: size (r), border irregularity (irr), heterogeneous pigment (het),
# asymmetric extra lobe (lobe) and how blurred the border is (soft); the ranges overlap on purpose
params = {0: dict(r=(4.0, 7.0), irr=(0.0, 0.10), het=(0.0, 0.06), lobe=0.05, soft=(1.2, 2.0)),
1: dict(r=(5.0, 8.5), irr=(0.07, 0.20), het=(0.04, 0.12), lobe=0.35, soft=(0.8, 1.6)),
2: dict(r=(6.0, 10.0), irr=(0.15, 0.32), het=(0.10, 0.22), lobe=0.75, soft=(0.5, 1.2))}
for i, c in enumerate(labels):
p = params[c]
cx, cy = rng.uniform(11, 17, 2)
r0, irr = rng.uniform(*p['r']), rng.uniform(*p['irr'])
het, soft = rng.uniform(*p['het']), rng.uniform(*p['soft'])
theta, dist = np.arctan2(yy - cy, xx - cx), np.hypot(xx - cx, yy - cy)
radius = r0 * (1 + sum(irr / k ** 0.5 * rng.uniform(0.5, 1.0) * np.cos(k * theta + rng.uniform(0, 2 * np.pi))
for k in range(2, 7)))
mask = 1 / (1 + np.exp(-(radius - dist) / soft))
if rng.random() < p['lobe']:
a = rng.uniform(0, 2 * np.pi)
lx, ly = cx + 0.8 * r0 * np.cos(a), cy + 0.8 * r0 * np.sin(a)
mask = np.maximum(mask, 1 / (1 + np.exp(-(0.55 * r0 - np.hypot(xx - lx, yy - ly)) / soft)))
pigment = rng.uniform(0.45, 0.6) + het * rng.normal(0, 1, (28, 28))
for _ in range(3): # darker and lighter blotches
bx, by = cx + rng.normal(0, r0 / 2), cy + rng.normal(0, r0 / 2)
pigment += het * 2.5 * rng.choice([-1, 1]) * np.exp(-((xx - bx) ** 2 + (yy - by) ** 2) / (2 * (r0 / 3) ** 2))
skin = 0.82 + 0.04 * (xx - 14) / 14 * rng.normal() + 0.03 * rng.normal(0, 1, (28, 28))
images[i] = np.clip(skin - mask * pigment, 0, 1)
return images, labels
def capture(images, rng):
"""Minermont's archive: each photo with its own exposure, noise and focus."""
out = np.empty_like(images)
for i, img in enumerate(images):
x = (img - 0.5) * rng.uniform(0.75, 1.25) + 0.5 + rng.uniform(-0.12, 0.12) # contrast and brightness
if rng.random() < 0.4: # out of focus: 3 x 3 blur
p = np.pad(x, 1, mode='edge')
x = sum(p[a:a + 28, b:b + 28] for a in range(3) for b in range(3)) / 9
out[i] = np.clip(x + rng.normal(0, rng.uniform(0.01, 0.06), x.shape), 0, 1) # sensor noise
return out
def archive(seed, counts):
"""Exactly counts[c] images of each class."""
images, labels = make_lesions(int(sum(counts) * 3.5), seed=seed)
idx = np.concatenate([np.where(labels == c)[0][:n] for c, n in enumerate(counts)])
return images[idx], labels[idx]
class_names = ['benign', 'suspicious', 'malignant']
clean, y_small = archive(1, (180, 70, 30))
X_small = capture(clean, np.random.default_rng(11))
X_more, y_more = archive(2, (180, 70, 30))
X_more = capture(X_more, np.random.default_rng(12))
X_test, y_test = make_lesions(900, seed=3)
X_test = capture(X_test, np.random.default_rng(13))
print("Training images per class:", dict(zip(class_names, np.bincount(y_small).tolist())))
print("Test images per class: ", dict(zip(class_names, np.bincount(y_test).tolist())))
fig, axes = plt.subplots(2, 6, figsize=(12, 4.4))
examples = [np.where(y_small == c)[0][k] for c in range(3) for k in range(2)]
for col, idx in enumerate(examples):
axes[0, col].imshow(clean[idx], cmap='gray', vmin=0, vmax=1)
axes[1, col].imshow(X_small[idx], cmap='gray', vmin=0, vmax=1)
axes[0, col].set_title(class_names[y_small[idx]])
for ax in axes.ravel():
ax.axis('off')
fig.suptitle('The same lesions before (top) and after the capture in the archive (bottom)'); plt.show()
Training images per class: {'benign': 180, 'suspicious': 70, 'malignant': 30}
Test images per class: {'benign': 460, 'suspicious': 276, 'malignant': 164}
Step 3: The network and how we measure it¶
The book's CNN: two convolutions of 16 and 32 filters with max-pooling, a dense layer of 64 neurons and 3 outputs, with an optional dropout, which switches off at random a share of the neurons during training so that the network cannot rely on any single one. train has an option for each remedy of the book: weights for the classes, augmentation, and which parameters are trained. Adam, 40 epochs, batches of 32.
With 280 images, the result of one training depends a lot on the random initialization, so each configuration is trained 5 times with different seeds, and evaluate reports the mean and the standard deviation of the accuracy and of the recall of each class on the 900 test images: the share of the lesions of that class that the network recognizes, the "85% on benign, 70% on suspicious, 65% on malignant" of Noah's first prototype.
def tensor(images):
return torch.tensor(images[:, None], dtype=torch.float32)
class CNN(nn.Module):
def __init__(self, dropout=0.0):
super().__init__()
self.features = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten())
self.classifier = nn.Sequential(nn.Dropout(dropout), nn.Linear(1568, 64), nn.ReLU(), nn.Dropout(dropout), nn.Linear(64, 3))
def forward(self, x):
return self.classifier(self.features(x))
def augment(x, g):
"""Random flips, quarter turns, brightness and contrast: a lesion has no 'up', and photos vary."""
if torch.rand(1, generator=g) < 0.5:
x = x.flip(3)
if torch.rand(1, generator=g) < 0.5:
x = x.flip(2)
x = torch.rot90(x, int(torch.randint(0, 4, (1,), generator=g)), (2, 3))
brightness = (torch.rand(len(x), 1, 1, 1, generator=g) - 0.5) * 0.1
contrast = 1 + (torch.rand(len(x), 1, 1, 1, generator=g) - 0.5) * 0.2
return ((x - 0.5) * contrast + 0.5 + brightness).clamp(0, 1)
def train(model, X, y, epochs=40, lr=1e-3, class_weights=None, augmentation=False, only_classifier=False, seed=0):
torch.manual_seed(seed)
g = torch.Generator().manual_seed(seed)
if only_classifier: # freeze the convolutions
for p in model.features.parameters():
p.requires_grad = False
optimizer = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=lr)
weights = None if class_weights is None else torch.tensor(class_weights, dtype=torch.float32)
Xt, yt = tensor(X), torch.tensor(y)
for _ in range(epochs):
model.train()
for batch in torch.randperm(len(yt), generator=g).split(32):
x = augment(Xt[batch], g) if augmentation else Xt[batch]
loss = nn.functional.cross_entropy(model(x), yt[batch], weight=weights)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return model
def scores(model):
model.eval()
with torch.no_grad():
pred = model(tensor(X_test)).argmax(dim=1).numpy()
return [np.mean(pred == y_test)] + [np.mean(pred[y_test == c] == c) for c in range(3)]
results = {}
def evaluate(name, make_model, seeds=5):
runs = []
for seed in range(seeds):
torch.manual_seed(seed) # the initial weights depend on the seed too
runs.append(scores(make_model(seed)))
runs = np.array(runs)
results[name] = runs.mean(axis=0)
spread = f" (±{runs[:, 0].std():.3f})" if seeds > 1 else ""
print(f"{name:<42} accuracy {runs[:, 0].mean():.3f}{spread:<9} | recall: benign {runs[:, 1].mean():.3f}, "
f"suspicious {runs[:, 2].mean():.3f}, malignant {runs[:, 3].mean():.3f} (±{runs[:, 3].std():.3f})")
Step 4: The starting point: training from scratch¶
The network starts from random weights and learns only from the 280 images.
evaluate('From scratch', lambda seed: train(CNN(), X_small, y_small, seed=seed))
From scratch accuracy 0.710 (±0.033) | recall: benign 0.927, suspicious 0.442, malignant 0.556 (±0.126)
The network learns the frequent class: it recognizes 0.927 of the benign lesions, but only 0.442 of the suspicious ones and 0.556 of the malignant ones, with a large spread from one training to another (±0.126): with 30 malignant images, whether the network learns them depends on luck. The overall accuracy, 0.710, hides it. It is the same pattern as Noah's first prototype, good on benign lesions and worst on malignant ones, "the ones we are most interested in", as Emma says.
Step 5: Class imbalance: class weights and oversampling¶
The two balancing techniques Emma mentions:
- Class weights, an adjustment of the loss: each error on an image of class $c$ counts $w_c = \frac{N}{3\,N_c}$ times, so the rare classes weigh as much in total as the frequent one. With 180, 70 and 30 images, $w = (0.52, 1.33, 3.11)$: a mistake on a malignant lesion costs six times as much as a mistake on a benign one.
- Oversampling: images of the rare classes are repeated, drawn at random with replacement, until the three classes have 180 each. Each epoch then has 540 images, so we train for 20 epochs to keep a similar number of updates.
class_weights = len(y_small) / (3 * np.bincount(y_small))
print("Class weights:", np.round(class_weights, 2))
def oversample(X, y, rng):
idx = np.concatenate([rng.choice(np.where(y == c)[0], np.bincount(y).max()) for c in range(3)])
return X[idx], y[idx]
evaluate('From scratch + class weights', lambda seed: train(CNN(), X_small, y_small, class_weights=class_weights, seed=seed))
evaluate('From scratch + oversampling', lambda seed: train(CNN(), *oversample(X_small, y_small, np.random.default_rng(seed)),
epochs=20, seed=seed))
Class weights: [0.52 1.33 3.11]
From scratch + class weights accuracy 0.746 (±0.009) | recall: benign 0.863, suspicious 0.584, malignant 0.690 (±0.047)
From scratch + oversampling accuracy 0.724 (±0.019) | recall: benign 0.769, suspicious 0.643, malignant 0.738 (±0.056)
Both techniques move the balance towards the rare classes. With class weights the recall of malignant lesions rises from 0.556 to 0.690 and that of suspicious ones from 0.442 to 0.584, while benign lesions fall from 0.927 to 0.863; the accuracy rises to 0.746 and the results are much more stable (±0.009). Oversampling goes further, 0.738 on malignant and 0.643 on suspicious lesions, at the price of more errors on benign ones, 0.769. Neither creates information: they only change what the network considers an expensive mistake.
Step 6: Data augmentation¶
Emma's other idea: "rotating, changing the lighting, or applying small distortions to the existing images". At every batch, augment flips the images, turns them by a random number of quarter turns (a lesion has no up or down) and changes their brightness and contrast a little, so the network never sees exactly the same image twice. We combine it with the class weights and a dropout of 0.3. First, eight versions of the same malignant lesion:
g = torch.Generator().manual_seed(0)
one = tensor(X_small[np.where(y_small == 2)[0][:1]])
fig, axes = plt.subplots(1, 8, figsize=(14, 2))
for ax in axes:
ax.imshow(augment(one, g)[0, 0], cmap='gray', vmin=0, vmax=1)
ax.axis('off')
plt.show()
evaluate('From scratch + weights + augmentation', lambda seed: train(CNN(dropout=0.3), X_small, y_small, class_weights=class_weights,
augmentation=True, seed=seed))
From scratch + weights + augmentation accuracy 0.755 (±0.017) | recall: benign 0.865, suspicious 0.562, malignant 0.770 (±0.025)
The eight versions are the same lesion turned, flipped and lit differently, and for the network they are eight different examples. With augmentation and dropout added to the class weights, the recall of malignant lesions rises to 0.770, with less spread (±0.025), and the accuracy to 0.755. It helps, but less than one might hope: augmentation varies what the network already has, it does not add new lesions.
Step 7: Transfer learning from another hospital's images¶
Transfer learning, in Emma's words, "consists of taking a model that has already been pre-trained on a large dataset and adapting it to our specific problem". Our large dataset is 6,000 labelled images from another hospital, the kind of images Grace proposed to use: the same kind of lesions, but taken with other equipment, brighter and with a different contrast, and with other proportions of the classes. We pretrain the network on them, and then:
- We use it as it is on Minermont's test images.
- We freeze its convolutions, the feature extractor, and retrain only the dense layers with the 280 images.
- We fine-tune the whole network with the 280 images, with a smaller learning rate.
X_other, y_other = make_lesions(6000, seed=4, shares=(0.4, 0.3, 0.3))
X_other = capture(np.clip(0.25 + 0.75 * X_other ** 1.4, 0, 1).astype(np.float32), np.random.default_rng(5)) # another camera
torch.manual_seed(0)
pretrained = train(CNN(dropout=0.3), X_other, y_other, epochs=8)
print(f"Pretrained on {len(y_other):,} images from another hospital\n")
evaluate('Pretrained elsewhere, used as it is', lambda seed: pretrained, seeds=1)
evaluate('Transfer: frozen convolutions', lambda seed: train(copy.deepcopy(pretrained), X_small, y_small, class_weights=class_weights,
augmentation=True, only_classifier=True, seed=seed))
evaluate('Transfer: whole network fine-tuned', lambda seed: train(copy.deepcopy(pretrained), X_small, y_small, lr=3e-4,
class_weights=class_weights, augmentation=True, seed=seed))
Pretrained on 6,000 images from another hospital Pretrained elsewhere, used as it is accuracy 0.741 | recall: benign 0.748, suspicious 0.601, malignant 0.957 (±0.000)
Transfer: frozen convolutions accuracy 0.813 (±0.007) | recall: benign 0.869, suspicious 0.678, malignant 0.883 (±0.009)
Transfer: whole network fine-tuned accuracy 0.824 (±0.010) | recall: benign 0.904, suspicious 0.658, malignant 0.880 (±0.007)
Used as it is, the network from the other hospital reaches an accuracy of 0.741, but it is not adapted to Minermont: it detects 0.957 of the malignant lesions and only 0.748 of the benign ones, because it learned with another camera and with more malignant lesions (30% of its images). Retraining only its dense layers with our 280 images gives 0.813, and fine-tuning the whole network 0.824, with a malignant recall of about 0.88 and a very small spread between trainings (±0.007 to ±0.010). The convolutions learned from 6,000 images are a better starting point than random weights, and our own images adapt them to our cameras and our proportions: exactly the two sides of Emma's argument.
Step 8: Twice as many images¶
Two months later the team has doubled the dataset. We simulate it with 280 more images taken like the first ones and repeat the best recipe without transfer and the recipe with transfer with 560 images, 360 benign, 140 suspicious and 60 malignant.
X_double, y_double = np.concatenate([X_small, X_more]), np.concatenate([y_small, y_more])
evaluate('560 images: from scratch + weights + augm.', lambda seed: train(CNN(dropout=0.3), X_double, y_double, class_weights=class_weights,
augmentation=True, seed=seed))
evaluate('560 images: transfer, fine-tuned', lambda seed: train(copy.deepcopy(pretrained), X_double, y_double, lr=3e-4,
class_weights=class_weights, augmentation=True, seed=seed))
560 images: from scratch + weights + augm. accuracy 0.763 (±0.027) | recall: benign 0.839, suspicious 0.598, malignant 0.830 (±0.043)
560 images: transfer, fine-tuned accuracy 0.827 (±0.003) | recall: benign 0.905, suspicious 0.662, malignant 0.884 (±0.007)
With 560 images, the network trained from scratch improves clearly: the recall of malignant lesions rises from 0.770 to 0.830 and the accuracy from 0.755 to 0.763. The transferred network barely changes, from 0.824 to 0.827 in accuracy and from 0.880 to 0.884 on malignant lesions: most of what it needed already came from the pretraining. Even with twice as many images, the network trained from scratch does not reach the one that started from the other hospital's images with half as many.
Step 9: Summary¶
names = list(results)
values = np.array([results[n] for n in names])
fig, ax = plt.subplots(figsize=(11, 5.5))
y_pos = np.arange(len(names))
ax.barh(y_pos - 0.2, values[:, 0], height=0.4, label='accuracy')
ax.barh(y_pos + 0.2, values[:, 3], height=0.4, label='recall of malignant lesions')
ax.set_yticks(y_pos, names); ax.invert_yaxis(); ax.set_xlim(0, 1); ax.grid(axis='x')
ax.legend(loc='upper right'); ax.set_title('Mean of 5 trainings on the 900 test images'); plt.tight_layout(); plt.show()
The orange bars tell the story that matters in dermatology: from 0.556 of the malignant lesions detected by the first network to about 0.88 with transfer learning. These are synthetic images and the numbers are ours, not the book's, but the remedies are the ones the team uses in the book.
Conclusion¶
- With few images, a network trained from scratch learns the frequent class and neglects the rare ones, which in medicine are usually the ones that matter.
- Class weights and oversampling tell the network that errors on the rare classes cost more. They raise the recall of malignant lesions at the price of some errors on benign ones.
- Data augmentation creates plausible variations of the images it already has, and together with dropout it makes the network depend less on the details of each photo.
- Transfer learning starts from features learned on a large dataset. Used as it is, a model from another hospital is not adapted to our images and our proportions, as Emma warns; adapted with our own 280 images, it is the best of the recipes.
- More data of our own helps most where the network has to learn everything from it: doubling the dataset improves the network trained from scratch much more than the transferred one.