Section 1 · Lesson 4

Model architecture

The project uses a U-Net: an encoder that shrinks the image into rich features, a decoder that grows them back to full size, and skip connections that carry fine detail across.

You'll learn

How a U-Net turns a 128 × 800 image into four mask channels.
What the skip connections add.
Which encoders the project tried, and why the input size is not resized.

Why a U-Net

A classifier shrinks an image to a few numbers. A segmentation model has to give an answer for every pixel, so it needs a way to get back to full resolution without losing the thin edges of a defect. A U-Net does this in two halves:

The encoder repeatedly halves the height and width while increasing the number of channels. Deep features see a large part of the image.
The decoder doubles the size again at each step.
At each size, a skip connection copies the encoder's features into the decoder, so the decoder can combine "what" (deep features) with "where" (sharp early features).

The model in the code

The project doesn't write the network by hand. src/train.py builds it with the segmentation_models_pytorch (smp) library:

def build_model(arch, encoder, pretrained=True):
    cls = getattr(smp, arch)
    return cls(encoder_name=encoder,
               encoder_weights="imagenet" if pretrained else None,
               classes=4)

The defaults are arch="Unet" and encoder="resnet34". The encoder starts from ImageNet-pretrained weights. The final layer has classes=4 output channels, one per defect type. The network outputs raw scores (logits), and a sigmoid turns each into a probability.

Why sigmoid and not softmax? An image can contain several defect types, and the four masks never need to add up to 1. A sigmoid per channel treats each class as its own yes-or-no question. This matches the label statistics from lesson 2: 1,198 images have two or more classes.

Data flow at the native size

The project feeds images at their native 128 × 800 size. Both numbers divide by 32, which a five-stage encoder needs, so the smallest feature map is 4 × 25. The milestone analysis also shows why not to resize to a square such as 256 × 256: that would stretch the width by 0.32 and the height by 2.0, which distorts the aspect ratio by a factor of 6.25.

Each image is scaled before it enters the network, in SegDataset.__getitem__:

img = torch.from_numpy(img.transpose(2, 0, 1).astype(np.float32) / 255.0)
img = (img - 0.5) / 0.25

U-Net explorer

Select a block to see its tensor shape for a 128 × 800 input. Turn the skip connections off to see which blocks lose their extra input. The shapes use the ResNet-34 encoder; EfficientNet-B3 has the same resolutions with different channel counts.

Skip connections

Encoders the project tried

Run	Encoder	Epochs	Learning rate	Notes
`r34_baseline`	`resnet34`	10	3e-4	Main model. Saved holdout probabilities, evaluation, and tuned post-processing.
`effb3`	`efficientnet-b3`	15 (log shows 8)	4e-4	Weights saved. The log records 8 epochs; no evaluation files exist for it.
`r34_pos`	`resnet34`	15	3e-4	Uses the positive-pairs Dice loss and 2× oversampling (lesson 5). Only the arguments file exists.

All three use arch="Unet", a batch size of 16, and seed 42, as recorded in each run's args.json. ResNet-34 is the baseline because each epoch takes roughly 2.5 to 3 times less time on the project's Apple-silicon machine: the logs show roughly 210 to 250 seconds per epoch for ResNet-34 and 535 to 730 seconds for EfficientNet-B3 (about 385 seconds in the last logged epoch).

What a convolution does

The encoder is built from convolutions: small kernels that slide over the image. Pick a 3 × 3 kernel or edit the numbers, and see the filtered version of a real training image. Deep networks learn thousands of such kernels instead of using hand-written ones.

Preset

Top: the input image. Bottom: the response after the kernel, scaled to the 0 to 255 range.

What is not in the model

The model has no special handling for empty images. It outputs a probability map for every class, and the post-processing step (lesson 6) decides whether a class counts as present. That separation lets the project tune the empty-image behaviour directly against the competition score.

What is the main job of a skip connection in a U-Net?

It passes high-resolution encoder features to the decoder so edges stay sharp.
It reduces the number of parameters.
It replaces the final sigmoid.

The decoder works at the deep layers' coarse resolution. The skip connection supplies the fine spatial detail from the matching encoder block.

Why does the model output four channels with a sigmoid on each?

Four channels use less memory.
Softmax can't be used with PyTorch.
An image can contain several defect types, so each class is a separate yes-or-no mask.

Softmax forces the classes at a pixel to compete and sum to 1. The labels allow several classes per image, so independent sigmoids fit better.