Building a U-Net
The shape that changed medical imaging.
In twenty fifteen, a small team at the University of Freiburg published a network for microscope images, trained on only thirty pictures. It became one of the most cited papers in the field, and its shape is still inside today's image generators. Let's build a tiny one.
- 2015: thirty training pictures
- One of the most cited papers
- Still inside image generators
The U shape
A classifier throws away the details about where things are. But segmentation needs to know exactly which pixel is which. So U-Net first shrinks the image to understand what's there. Then it grows back to full size. And here's the secret. Skip connections pass the fine details straight across, from the shrinking side to the growing side.
A tiny U-Net
Our data is simple shapes on a noisy background, with a mask for each. In the forward pass, the down layer keeps a copy called skip. Then we shrink, and upsample. The key line is torch dot cat, which joins the upsampled features to that skip copy. The output has one score per pixel. Run it. In six short epochs, on a laptop, the overlap on new images is around point nine. Every pixel gets its own answer. And you've built it in thirty lines.
import numpy as np
import cv2
import torch
from torch import nn
def sample(n, seed=0):
rng = np.random.default_rng(seed)
X, M = [], []
for _ in range(n):
img = rng.normal(90, 12, (64, 64, 1)).astype(np.float32)
mask = np.zeros((64, 64), np.uint8)
for _ in range(int(rng.integers(1, 4))):
at = tuple(int(v) for v in rng.integers(10, 54, 2))
cv2.circle(mask, at, int(rng.integers(5, 11)), 1, -1)
img[mask == 1] += 60
X.append(img / 255); M.append(mask)
return torch.tensor(np.array(X)).permute(0, 3, 1, 2), torch.tensor(np.array(M)).float()
class UNet(nn.Module):
def __init__(self):
super().__init__()
self.down = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU())
self.bottom = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
self.up = nn.ConvTranspose2d(32, 16, 2, stride=2)
self.out = nn.Conv2d(32, 1, 1)
def forward(self, x):
skip = self.down(x)
up = self.up(self.bottom(skip))
return self.out(torch.cat([up, skip], 1)).squeeze(1)
torch.manual_seed(0)
X, M = sample(600)
model = UNet()
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
for epoch in range(6):
for i in torch.randperm(len(X)).split(32):
loss = nn.functional.binary_cross_entropy_with_logits(model(X[i]), M[i])
opt.zero_grad(); loss.backward(); opt.step()
Xt, Mt = sample(100, seed=9)
p = model(Xt) > 0
print(f"Test IoU: {(p & (Mt > 0)).sum() / (p | (Mt > 0)).sum():.3f}")Test IoU: 0.896
Summary
Shrink, grow, and skip. That shape is everywhere. Next, how can you find a defect that you've never seen before?
Key takeaways- Skip connections keep detail
- One score per pixel
- It's in today's image generators