Section 1 · Lesson 5

Training the model

The recipe is short: BCE plus Dice loss, light augmentation, AdamW with a one-cycle schedule, and a holdout score after every epoch.

The recipe

The whole training loop is short enough to read in one go:

opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
steps = args.epochs * len(tr_dl)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.1)
bce = torch.nn.BCEWithLogitsLoss()
dice = pos_dice_loss if args.loss == "pos" else smp.losses.DiceLoss("multilabel", from_logits=True)

for img, msk in tr_dl:
    logits = model(img)
    loss = bce(logits, msk) + args.dice_w * dice(logits, msk)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    sched.step()

Loss

The loss adds two terms. Binary cross-entropy (BCE) looks at every pixel on its own and gives smooth gradients. Dice loss is \(1 - \mathrm{Dice}\) computed on the predicted probabilities, so it directly pushes up the overlap that the metric measures. It also copes with the fact that defect pixels are a tiny fraction of the image.

The project added a second option, --loss pos, for an experiment. It averages Dice only over pairs whose true mask is non-empty, so empty pairs don't pull the loss toward predicting nothing:

def pos_dice_loss(logits, targets):
    probs = logits.sigmoid()
    inter = (probs * targets).sum((-2, -1))
    denom = probs.sum((-2, -1)) + targets.sum((-2, -1))
    dice = (2 * inter + 1) / (denom + 1)
    pos = (targets.sum((-2, -1)) > 0).float()
    return ((1 - dice) * pos).sum() / pos.sum().clamp_min(1)

The r34_pos run used this loss with --oversample 2.0, which draws images containing defect_a or defect_b twice as often. The folder holds only its arguments file, and its log is empty, so this course can't report a result for it.

Augmentation

Augmentation shows the model slightly different versions of each image. The project applies these with the Albumentations library, to the image and the mask together:

A.Compose([
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.RandomBrightnessContrast(0.15, 0.15, p=0.5),
    A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=0,
                       border_mode=0, p=0.3),
])

The mask must receive the same geometric change as the image, otherwise the labels no longer line up. Brightness and contrast change only the image.

The learning-rate schedule

The one-cycle schedule starts at a low rate, ramps up to max_lr over the first 10% of the steps, and then decays smoothly toward zero. Warm-up protects the pretrained weights early on, and the long decay lets the model settle. PyTorch's defaults start at max_lr / 25 and end at that start value divided by 10,000.

Reading the curves

The training script prints one JSON line per epoch with the loss and the holdout score at four thresholds. Here are the real numbers for the baseline:

Four observations from the files:

• Epoch 2 is worse than epoch 1. The holdout score at tau 0.3 drops from 0.4309 to 0.4141, and foreground Dice from 0.2825 to 0.1656. The learning rate peaks around the end of epoch 1, so a brief setback while it is highest is plausible. The model recovers by epoch 3.
• The score is still rising at epoch 10. The score at 0.3 goes 0.6186, 0.6205, 0.6271, 0.6284 over the last four epochs, so a longer run might help. The files don't test that.
• The EfficientNet-B3 curve is noisy. Its score at tau 0.5 reaches 0.6503 at epoch 6, falls to 0.5928 at epoch 7, and recovers to 0.6402 at epoch 8. A holdout with few images of the rare classes can make the score swing between epochs.
• Flip augmentation at test time helps. After the last epoch, averaging the model's predictions over the original and two flipped copies lifts the score at 0.3 from 0.6284 to 0.6420 (log line "TTA"). Lesson 7 explains this.

The Milestone 1 starter model, trained for 8 epochs, reports a final validation loss of 0.2785 and an assessment score of 0.5627 to 0.5677 across thresholds. It is a reference point, not a rival: its numbers come from a different script and loss.