new build model

The model returned cells as [row, column] while the dataset targets,
convert_prediction and the metrics use [column, row], so most cells
were trained on faces outside the region they can see.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

api/analytics/run_viewer.html

@@ -396,6 +396,7 @@ function fmtParams(n) {
   return String(n);
 }
 const metricLabel = (m) => m.replace(/_/g, " ");
+const extraMetricNames = (runs) => [...new Set([...runs].flatMap((r) => Object.keys(r.summary.atBest)))];
 
 // ------------------------------------------------------------------ loading
 
@@ -497,6 +498,13 @@ function makeRun(dir, sourceId, info, metrics, architecture, diff) {
   const valLoss = col("val_loss"), valAcc = col("val_acc");
   const last = rows.length ? rows[rows.length - 1] : {};
   const result = info.result || {};
+  // Extra val metrics (e.g. YOLO's conf_ap): the value at the best epoch, whose weights were kept.
+  const bestRow = rows.find((r) => r.epoch === result.best_epoch) || last;
+  const atBest = {};
+  for (const c of metrics.columns) {
+    const m = c.match(/^val_(.+)$/);
+    if (m && m[1] !== "loss" && m[1] !== "acc" && typeof bestRow[c] === "number") atBest[m[1]] = bestRow[c];
+  }
   return {
     key: dir,
     source: sourceId,
@@ -519,6 +527,7 @@ function makeRun(dir, sourceId, info, metrics, architecture, diff) {
       finalTrainLoss: typeof last.train_loss === "number" ? last.train_loss : null,
       bestEpoch: result.best_epoch ?? null,
       monitor: result.monitor || "loss",
+      atBest,
       duration: result.duration_s ?? null,
       params: (info.model || {}).params ?? null,
       archHash: (info.model || {}).architecture_hash || null,
@@ -676,6 +685,10 @@ function runsColumns(varying) {
     { id: "minValLoss", label: "Min val loss", get: (r) => r.summary.minValLoss, num: true },
   ];
   if (hasAcc) cols.push({ id: "maxValAcc", label: "Max val acc", get: (r) => r.summary.maxValAcc, num: true });
+  for (const name of extraMetricNames(state.runs.values())) {
+    cols.push({ id: "best:" + name, label: metricLabel(name), title: `val ${metricLabel(name)} at the best epoch`,
+      get: (r) => r.summary.atBest[name], num: true });
+  }
   cols.push(
     { id: "bestEpoch", label: "Best epoch", get: (r) => r.summary.bestEpoch, num: true },
     { id: "duration", label: "Duration", get: (r) => r.summary.duration, num: true },
@@ -728,7 +741,7 @@ function renderRunsTable() {
     const th = h("th", {
       class: [sortable ? "sortable" : "", c.num ? "num" : "", c.cfg ? "cfg" : ""].join(" ").trim(),
       scope: "col",
-      title: c.cfg ? "setting that differs between the listed runs" : null,
+      title: c.title || (c.cfg ? "setting that differs between the listed runs" : null),
       "aria-sort": state.sort.col === c.id ? (state.sort.dir > 0 ? "ascending" : "descending") : null,
     }, c.label, arrow ? h("span", { class: "arrow" }, arrow) : null);
     if (sortable) th.addEventListener("click", () => {
@@ -1017,13 +1030,16 @@ function renderScatter() {
   const prevX = xSel.value || recall("scatterX");
   xSel.replaceChildren(...xOptions.map((o) => h("option", { value: o.id }, o.label)));
   xSel.value = xOptions.some((o) => o.id === prevX) ? prevX : xOptions[0].id;
-  const yOptions = Y_MEASURES.filter((mm) => runs.some((r) => mm.get(r) != null));
+  const measures = [...Y_MEASURES.slice(0, 2),
+    ...extraMetricNames(runs).map((name) => ({ id: "best:" + name, label: `val ${metricLabel(name)} at best epoch`, get: (r) => r.summary.atBest[name] })),
+    ...Y_MEASURES.slice(2)];
+  const yOptions = measures.filter((mm) => runs.some((r) => mm.get(r) != null));
   const prevY = ySel.value || recall("scatterY");
   ySel.replaceChildren(...yOptions.map((o) => h("option", { value: o.id }, o.label)));
   ySel.value = yOptions.some((o) => o.id === prevY) ? prevY : (yOptions[0] || {}).id || "";
 
   const xo = xOptions.find((o) => o.id === xSel.value);
-  const yo = Y_MEASURES.find((o) => o.id === ySel.value);
+  const yo = measures.find((o) => o.id === ySel.value);
   const container = $("s-chart");
   if (!xo || !yo) { container.replaceChildren(h("div", { class: "empty" }, "Nothing to plot.")); return; }
 
@@ -1131,6 +1147,7 @@ function renderCompare() {
     const f = flatten({ ...r.info, git: { ...(r.info.git || {}), changed_files: ((r.info.git || {}).changed_files || []).length } }, "", {});
     f["summary.min_val_loss"] = r.summary.minValLoss;
     if (r.summary.maxValAcc != null) f["summary.max_val_acc"] = r.summary.maxValAcc;
+    for (const [name, v] of Object.entries(r.summary.atBest)) f[`summary.val_${name}_at_best_epoch`] = v;
     f["summary.duration"] = fmtDuration(r.summary.duration);
     return f;
   });

api/doc/yolo.md

@@ -1,325 +0,0 @@
-# YOLO face detector
-
-A small single-shot face detector in the style of YOLO v1. It splits the image
-into a 6 × 6 grid and predicts at most one face per grid cell. The website uses
-it for live webcam face detection over a websocket.
-
-| | |
-|---|---|
-| Model code | [`api/src/models/yolo.py`](../src/models/yolo.py) (`Yolo_model`, `SkipBlock`, `convert_prediction`) |
-| Training code | [`api/src/yolo/`](../src/yolo/) (`config.py`, `dataset.py`, `loss.py`, `tasks.py`) |
-| Deployed weights | `api/build_models/face_detection_yolo` (PyTorch `state_dict`) |
-| Production use | [`api/src/production/inference.py`](../src/production/inference.py) `detect_faces`, websocket `/predict_face` |
-| Input | RGB image resized to 64 × 64, values in [0, 1] |
-| Output | 6 × 6 grid × 6 values per cell (x, y, w, h, confidence, class) |
-| Classes | 1 (face) |
-| Parameters | 657,254 (with `c_hidden=16`) |
-| Speed | about 18 ms per image on CPU, batch size 1 (measured on the dev machine) |
-
-## Architecture
-
-The network is a fully convolutional stack of residual "skip blocks". It never
-downsamples: every layer works at the full 64 × 64 resolution. At the end, one
-`AdaptiveAvgPool2d` reduces the feature map to the 6 × 6 grid, and a 1 × 1
-convolution turns each grid cell's features into a prediction.
-
-```mermaid
-flowchart TD
-    IN["Input image<br/>3 × 64 × 64"]
-
-    subgraph STEM["Stem"]
-        S1["Conv 3×3, 3 → 16<br/>GroupNorm(2) · LeakyReLU"]
-    end
-
-    subgraph ST1["Stage 1 · 16 channels"]
-        B1["SkipBlock 16 → 16"] --> B2["SkipBlock 16 → 16"] --> B3["SkipBlock 16 → 16"] --> B4["SkipBlock 16 → 16"]
-    end
-
-    subgraph ST2["Stage 2 · 32 channels"]
-        B5["SkipBlock 16 → 32"] --> B6["SkipBlock 32 → 32"] --> B7["SkipBlock 32 → 32"] --> B8["SkipBlock 32 → 32"]
-    end
-
-    subgraph ST3["Stage 3 · 64 channels"]
-        B9["SkipBlock 32 → 64"] --> B10["SkipBlock 64 → 64"] --> B11["SkipBlock 64 → 64"] --> B12["SkipBlock 64 → 64"]
-    end
-
-    subgraph HEAD["Head"]
-        H1["Conv 3×3, 64 → 128<br/>GroupNorm(8) · LeakyReLU · Dropout 0.3"]
-        H2["AdaptiveAvgPool2d → 6 × 6<br/>128 × 6 × 6"]
-        H3["Conv 1×1, 128 → 6<br/>6 × 6 × 6"]
-        H1 --> H2 --> H3
-    end
-
-    ACT["permute to [B, 6, 6, 6]<br/>activation per channel"]
-    OUT["Per cell: x · y · w · h · conf · class"]
-
-    IN --> STEM --> ST1 --> ST2 --> ST3 --> HEAD --> ACT --> OUT
-```
-
-Every layer up to and including the head's 3 × 3 conv keeps the spatial size at
-64 × 64. Only the channel count changes.
-
-| Index in `model` | Layer | Output shape (batch 1) | Parameters |
-|---|---|---|---|
-| 0–2 | Conv 3×3 + GroupNorm + LeakyReLU | 16 × 64 × 64 | 480 |
-| 3–6 | 4 × SkipBlock 16 → 16 | 16 × 64 × 64 | 4 × 7,360 |
-| 7–10 | SkipBlock 16 → 32, 3 × SkipBlock 32 → 32 | 32 × 64 × 64 | 23,936 + 3 × 29,056 |
-| 11–14 | SkipBlock 32 → 64, 3 × SkipBlock 64 → 64 | 64 × 64 × 64 | 94,976 + 3 × 115,456 |
-| 15–18 | Conv 3×3 + GroupNorm + LeakyReLU + Dropout | 128 × 64 × 64 | 74,112 |
-| 19 | AdaptiveAvgPool2d(6, 6) | 128 × 6 × 6 | 0 |
-| 20 | Conv 1×1 | 6 × 6 × 6 | 774 |
-
-The channel counts scale with `c_hidden` (default 16): c, 2c, 4c, then 8c in the
-head. The output has `boxes * 5 + labels` channels. With 1 box and 1 label
-that is 6.
-
-### SkipBlock
-
-Each block is a residual unit. Three 3 × 3 convolutions form the main path and
-a 1 × 1 convolution forms the shortcut. The 1 × 1 conv lets the shortcut change
-the channel count.
-
-```mermaid
-flowchart LR
-    X["x<br/>C_in × 64 × 64"]
-
-    subgraph MAIN["main path"]
-        C1["Conv 3×3 → C_out<br/>GroupNorm · LeakyReLU"]
-        C2["Conv 3×3<br/>GroupNorm · LeakyReLU"]
-        C3["Conv 3×3<br/>GroupNorm · LeakyReLU"]
-        C1 --> C2 --> C3
-    end
-
-    subgraph SKIP["shortcut"]
-        K["Conv 1×1 → C_out<br/>GroupNorm · LeakyReLU"]
-    end
-
-    ADD(("+"))
-    POST["LeakyReLU<br/>Dropout p=0.3 (always on)"]
-    Y["y<br/>C_out × 64 × 64"]
-
-    X --> C1
-    X --> K
-    C3 --> ADD
-    K --> ADD
-    ADD --> POST --> Y
-```
-
-All GroupNorm layers use `C / 8` groups, so each group holds 8 channels.
-
-### Design notes
-
-- **No downsampling.** All twelve blocks run at 64 × 64. This makes the network
-  cheap in parameters but comparatively expensive in compute and memory for its
-  size.
-- **Receptive field.** The network has 38 stacked 3 × 3 convolutions (1 stem,
-  36 in the blocks, 1 in the head). That gives a receptive field of
-  1 + 38 × 2 = 77 px, so every output cell can see the whole 64 px image.
-- **Uneven pooling bins.** 64 is not divisible by 6, so `AdaptiveAvgPool2d`
-  averages bins of 10–11 px, and neighbouring bins overlap slightly.
-- **Layer names are part of the file format.** `build_models/face_detection_yolo`
-  is a plain `state_dict` keyed by names like `model.3.conv.0.weight`. Adding,
-  removing or reordering layers in `Yolo_model` breaks loading of the deployed
-  weights.
-
-## Output encoding
-
-`forward` returns a tensor of shape `[batch, 6, 6, 6]`. For each grid cell
-`[i, j]` the last dimension holds:
-
-| Channel | Meaning | Range | Activation |
-|---|---|---|---|
-| 0 | `x`: box centre inside the cell, as a fraction of the cell width | 0–1 | sigmoid |
-| 1 | `y`: box centre inside the cell, as a fraction of the cell height | 0–1 | sigmoid |
-| 2 | `w`: box width as a fraction of the image width | ≥ 0 | `exp` or sigmoid |
-| 3 | `h`: box height as a fraction of the image height | ≥ 0 | `exp` or sigmoid |
-| 4 | confidence that a face's centre lies in this cell | 0–1 | sigmoid |
-| 5 | class score (face) | 0–1 | sigmoid |
-
-`size_activation` decides how `w` and `h` are produced:
-
-- `"exp"`: `w = exp(raw)`, unbounded. The **deployed** weights were trained
-  this way, so production sets `FACE_SIZE_ACTIVATION = "exp"`.
-- `"sigmoid"`: `w = sigmoid(raw)`, bounded to 0–1. This is the **training
-  default** in `Config.size_activation`.
-
-A model must always run with the activation it was trained with.
-
-**Grid index convention.** The first grid index `i` is the horizontal cell
-(column) and the second `j` is the vertical cell (row). The dataset builds
-targets this way, and `convert_prediction` reads them back the same way.
-
-The convolutional feature map is laid out `[row, column]`, so the network
-effectively learns a spatial transpose. Its receptive field covers the whole
-image, so it can do that. A check on the deployed model confirms it: on 69
-validation images with one face in an off-diagonal cell, the most confident
-cell matched the `[column, row]` target 49 times and the transposed cell once.
-If you change the dataset or the decoding, keep both on the same convention.
-
-## From prediction to boxes
-
-`convert_prediction(label, image, threshold)` turns one `[6, 6, 6]` output into
-pixel boxes. It also works on a target tensor, which is how `view-data` draws
-the ground truth.
-
-```mermaid
-flowchart LR
-    P["Prediction<br/>6 × 6 × 6"] --> F{"conf ≥ threshold?"}
-    F -- no --> E["empty cell<br/>(only used for drawing)"]
-    F -- yes --> D["decode cell i, j<br/>to centre + size in px"]
-    D --> C["corners<br/>xmin, ymin, xmax, ymax"]
-    C --> N["NMS<br/>IoU threshold 0.1"]
-    N --> S["scale to the original<br/>image size (production)"]
-```
-
-For a cell `[i, j]` on a 64 px image with a 6 × 6 grid (cell size
-`c = 64 / 6 ≈ 10.67 px`):
-
-```
-centre_x = (i + x) · c          width  = w · 64
-centre_y = (j + y) · c          height = h · 64
-
-xmin = centre_x − width / 2     xmax = xmin + width     (width and height are at least 1 px)
-ymin = centre_y − height / 2    ymax = ymin + height
-```
-
-Non-maximum suppression (`nms` in
-[`api/src/common/boxes.py`](../src/common/boxes.py)) sorts the boxes by
-confidence. It keeps the best box and drops every other box that overlaps it
-with IoU ≥ 0.1. The low threshold suits faces, which rarely overlap.
-
-In production, `detect_faces` multiplies the boxes by `W / 64` and `H / 64` to
-map them back to the size of the uploaded frame.
-
-## Training
-
-### Data
-
-- **Format.** COCO: images plus `_annotations.coco.json` in each folder.
-  Defaults: `api/data/faces_scenery/train` and `api/data/faces_scenery/test`.
-- **Resize.** Images are resized to 64 × 64.
-- **Augmentation** (albumentations):
-  - horizontal flip, p = 0.5
-  - brightness/contrast, p = 0.2
-  - affine: translate ±40 %, scale 0.5–1.5, edge pixels replicated, p = 0.5
-  - boxes that keep less than 50 % of their area are dropped
-- **Augmented validation.** `augment_val=True` by default, so the validation
-  set is augmented too. The original code did this.
-- **Targets.** `YoloDataset` builds a `[6, 6, 5 + classes]` target per image.
-  - Each cell takes the largest face whose centre lies in it. Annotations are
-    sorted by area, largest first, and each face is assigned only once.
-  - `x` and `y` are the centre's offset inside the cell, divided by the cell
-    size.
-  - `w` and `h` are the box size divided by the image size.
-  - Confidence is 1, and the class is one-hot.
-  - All other cells are zero.
-
-### Loss
-
-`YoloLoss` ([`api/src/yolo/loss.py`](../src/yolo/loss.py)) is the sum of
-four unweighted terms:
-
-| Term | Cells | Function |
-|---|---|---|
-| box | cells with a face | MSE on (x, y, w, h) |
-| object | cells with a face | MSE on confidence vs 1 |
-| no-object | empty cells | MSE on confidence vs 0 |
-| class | cells with a face | binary cross-entropy (with logits) |
-
-Each term is averaged over its own set of cells. So the roughly 35 empty cells
-together weigh the same as the one or two face cells, without the λ factors the
-original YOLO paper uses.
-
-### Running
-
-The commands run from `api/`. Every `Config` field can be overridden on the
-command line.
-
-```bash
-python -m src.main yolo.train                    # train (500 epochs, batch 64, Adam lr 1e-3, wd 1e-4)
-python -m src.main yolo.train --epochs 50 --size-activation exp
-python -m src.main yolo.view-data                # show validation images with their target grid
-python -m src.main yolo.sample                   # show the trained model's predictions
-python -m src.main yolo.webcam --camera 0        # live detection on a local webcam, q quits
-python -m src.main yolo.benchmark-workers        # find the fastest DataLoader num_workers
-python -m src.main yolo.train --help             # list every option
-```
-
-- **Checkpoints.** Training keeps the checkpoint with the lowest validation
-  loss at `saved_models/face_detection_yolo/face_detection_yolo`.
-- **Drawn grids.** `sample`, `view-data` and `webcam` draw red boxes for
-  detections, blue for cells with a face, and green for empty cells.
-- **Thresholds.** These tasks draw a box from confidence 0.95 (`--threshold`).
-  Production uses 0.9.
-
-### Deploying a new model
-
-1. Train it; the result is in `api/saved_models/face_detection_yolo/`.
-2. Copy that file to `api/build_models/face_detection_yolo`.
-3. Make the settings in
-   [`api/src/production/settings.py`](../src/production/settings.py) match
-   how it was trained:
-   - `FACE_SIZE_ACTIVATION`: `"sigmoid"` if you kept the training default.
-   - `FACE_C_HIDDEN`, `FACE_GRID` and `FACE_IMAGE_SIZE`.
-4. Restart the API. A size mismatch fails at startup in `load_checkpoint`.
-   An activation mismatch does not fail; it silently produces wrong box sizes.
-
-## Use in production
-
-```mermaid
-sequenceDiagram
-    participant B as Browser (Object_Detection.jsx)
-    participant W as /predict_face websocket
-    participant I as detect_faces
-    participant M as Yolo_model
-
-    B->>W: webcam frame as JPEG bytes
-    Note over W: checks origin, connection limits,<br/>frame size and FACE_MAX_FPS.<br/>Frames that arrive too fast or<br/>while busy are dropped
-    W->>I: frame bytes
-    I->>I: decode, resize to 64 × 64
-    I->>M: tensor [1, 3, 64, 64]
-    M-->>I: [1, 6, 6, 6]
-    I->>I: convert_prediction (threshold 0.9) + NMS,<br/>scale to frame size
-    I-->>W: boxes
-    W-->>B: {"bboxes": [[xmin, ymin, xmax, ymax], ...]}
-    B->>B: draw boxes on the overlay canvas
-```
-
-Relevant settings in `api/src/production/settings.py`:
-
-| Setting | Value | Meaning |
-|---|---|---|
-| `FACE_MODEL_FILE` | `face_detection_yolo` | file in `api/build_models/` |
-| `FACE_IMAGE_SIZE` | 64 | model input size |
-| `FACE_GRID` | 6 | grid size |
-| `FACE_C_HIDDEN` | 16 | base channel count |
-| `FACE_SIZE_ACTIVATION` | `"exp"` | must match the weights |
-| `FACE_THRESHOLD` | 0.9 | minimum confidence for a box |
-| `FACE_MAX_FPS` | 5.5 (env) | frames per second processed per websocket |
-
-## Known quirks and limitations
-
-- **Dropout is always on.** `SkipBlock.forward` calls `F.dropout(..., p=0.3)`,
-  which defaults to `training=True`, so it stays active in `eval()` mode.
-  - Two runs on the same image give slightly different outputs, and boxes can
-    flicker between video frames.
-  - The deployed weights were trained this way, so it is left as is.
-  - The head's `nn.Dropout` behaves normally and is off in eval mode.
-- **Activation default differs from the deployed model.** Training defaults to
-  `"sigmoid"`, while the deployed weights need `"exp"`. See
-  [Deploying a new model](#deploying-a-new-model).
-- **Class loss applies a sigmoid twice.** The model already outputs sigmoid
-  values, and `binary_cross_entropy_with_logits` applies another sigmoid. The
-  class output is unused with a single class (`convert_prediction` ignores it),
-  so this does not affect detections.
-- **One face per cell, at most 36 faces.** When two face centres fall into the
-  same cell, only the larger one is a training target, and the model can
-  predict only one box there.
-- **No anchors, one box per cell.** `boxes > 1` is accepted by the constructor,
-  but the dataset, loss and `convert_prediction` only handle the first box.
-- **Small input.** At 64 × 64, faces that are a few pixels tall after resizing
-  are hard to detect. Very wide or tall frames are squashed, because the
-  resize does not keep the aspect ratio.
-- **Noisy validation loss.** The validation set is augmented by default, so the
-  loss that picks the best checkpoint changes from epoch to epoch even with
-  fixed weights.

api/src/common/boxes.py

@@ -54,10 +54,13 @@ def image_to_cell_coords(x, y, img_w, img_h, S, cell_i, cell_j):
     return x - (img_w / S) * cell_i, y - (img_h / S) * cell_j
 
 
+def grid_cell(x, y, img_w, img_h, S):
+    """Image position -> (cell_i, cell_j) of the S x S grid cell it falls in."""
+    return min(int(x / (img_w / S)), S - 1), min(int(y / (img_h / S)), S - 1)
+
+
 def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
-    gt_cell_i = min(int(x / (img_w / S)), S - 1)
-    gt_cell_j = min(int(y / (img_h / S)), S - 1)
-    return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
+    return grid_cell(x, y, img_w, img_h, S) == (cell_i, cell_j)
 
 
 def scale_box(box, scale_w, scale_h):

api/src/common/training.py

@@ -4,6 +4,9 @@ Each epoch trains, validates, prints the metrics and saves the model when the
 monitored metric improves. `prepare_batch` turns a loader batch into
 (inputs, targets); the default just moves both to the device.
 
+An optional `scheduler` is stepped once after every epoch, so the learning
+rate logged for an epoch is the one it trained with.
+
 Passing the task's `config` also records the run (settings, architecture,
 metrics per epoch, best weights) under config.runs_dir, see run_log.py."""
 
@@ -23,34 +26,72 @@ def _default_prepare(batch, device):
     return inputs.to(device), targets.to(device)
 
 
-def _run_epoch(model, loss_fn, loader, device, prepare_batch, track_accuracy, optimizer=None, desc=""):
+def adamw(model, lr, weight_decay):
+    """AdamW that decays only conv/linear weights. Biases and BatchNorm
+    parameters (everything 1-D) are left alone, shrinking them only hurts."""
+    decay = [p for p in model.parameters() if p.requires_grad and p.ndim > 1]
+    no_decay = [p for p in model.parameters() if p.requires_grad and p.ndim <= 1]
+    return torch.optim.AdamW([{"params": decay, "weight_decay": weight_decay},
+                              {"params": no_decay, "weight_decay": 0.0}], lr=lr)
+
+
+def warmup_cosine(optimizer, epochs, warmup_epochs, min_lr):
+    """Per-epoch schedule: linear warmup from lr/10 over warmup_epochs, then
+    cosine decay from lr to min_lr over the remaining epochs."""
+    floor = min_lr / optimizer.param_groups[0]["lr"]
+    decay_epochs = max(epochs - warmup_epochs - 1, 1)
+
+    def factor(epoch):  # epoch counts from 0
+        if epoch < warmup_epochs:
+            return 0.1 + 0.9 * epoch / warmup_epochs
+        progress = min(epoch - warmup_epochs, decay_epochs) / decay_epochs
+        return floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * progress))
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, factor)
+
+
+def _run_epoch(model, loss_fn, loader, device, prepare_batch, track_accuracy, optimizer=None,
+               val_metrics=None, desc=""):
     training = optimizer is not None
     model.train(training)
 
     loss_sum, correct, count = 0.0, 0, 0
+    all_preds, all_targets = [], []
     with torch.set_grad_enabled(training):
         for batch in tqdm(loader, desc=desc, leave=False):
             inputs, targets = prepare_batch(batch, device)
 
             preds = model(inputs)
-            loss = loss_fn(preds, targets)
 
             if training:
+                loss = loss_fn(preds, targets)
                 optimizer.zero_grad()
                 loss.backward()
                 optimizer.step()
+                loss_sum += loss.item() * inputs.size(0)
+            else:
+                all_preds.append(preds)
+                all_targets.append(targets)
 
-            loss_sum += loss.item() * inputs.size(0)
             count += inputs.size(0)
             if track_accuracy:
                 correct += (preds.argmax(dim=1) == targets).sum().item()
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    metrics = {"loss": loss_sum / max(count, 1)}
+    if training:
+        metrics = {"loss": loss_sum / max(count, 1)}
+    else:
+        # Validation loss over the whole set at once, so it doesn't depend on
+        # the batch size (YoloLoss averages per batch over the cells with a face).
+        preds, targets = torch.cat(all_preds), torch.cat(all_targets)
+        metrics = {"loss": loss_fn(preds, targets).item()}
+        if val_metrics:
+            metrics.update(val_metrics(preds, targets))
+        del preds, targets, all_preds, all_targets
     if track_accuracy:
         metrics["acc"] = correct / max(count, 1)
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     return metrics
 
 
@@ -65,10 +106,12 @@ def fit(
     save_dir: str,
     model_name: str,
     save: bool = True,
-    monitor: str = "loss",          # "loss" (lower is better) or "acc" (higher is better)
+    monitor: str = "loss",          # "loss" (lower is better) or any other val metric (higher is better)
     track_accuracy: bool = False,   # for classifiers
     prepare_batch=_default_prepare,
+    val_metrics=None,               # f(preds, targets) -> dict of extra metrics over the whole val set
     config=None,                    # the task's Config; enables the run log
+    scheduler=None,                 # LR scheduler, stepped once per epoch
 ):
     run = None
     if config is not None and config.runs_dir:
@@ -76,7 +119,8 @@ def fit(
                      train_loader, val_loader, device, monitor)
     try:
         best = _fit_epochs(model, loss_fn, optimizer, train_loader, val_loader, device, epochs,
-                           save_dir, model_name, save, monitor, track_accuracy, prepare_batch, run)
+                           save_dir, model_name, save, monitor, track_accuracy, prepare_batch, val_metrics, run,
+                           scheduler)
     except KeyboardInterrupt:
         if run:
             run.finish("interrupted")
@@ -91,7 +135,8 @@ def fit(
 
 
 def _fit_epochs(model, loss_fn, optimizer, train_loader, val_loader, device, epochs,
-                save_dir, model_name, save, monitor, track_accuracy, prepare_batch, run):
+                save_dir, model_name, save, monitor, track_accuracy, prepare_batch, val_metrics, run,
+                scheduler):
     best = math.inf if monitor == "loss" else -math.inf
 
     for epoch in range(1, epochs + 1):
@@ -99,7 +144,7 @@ def _fit_epochs(model, loss_fn, optimizer, train_loader, val_loader, device, epo
         train = _run_epoch(model, loss_fn, train_loader, device, prepare_batch, track_accuracy,
                            optimizer=optimizer, desc=f"Train {epoch}")
         val = _run_epoch(model, loss_fn, val_loader, device, prepare_batch, track_accuracy,
-                         desc=f"Validate {epoch}")
+                         val_metrics=val_metrics, desc=f"Validate {epoch}")
 
         improved = val[monitor] < best if monitor == "loss" else val[monitor] > best
         if improved:
@@ -111,11 +156,14 @@ def _fit_epochs(model, loss_fn, optimizer, train_loader, val_loader, device, epo
         if run:
             run.log_epoch(epoch, train, val, lr=optimizer.param_groups[0]["lr"],
                           seconds=time.perf_counter() - start, improved=improved)
+        if scheduler:
+            scheduler.step()
 
         parts = [f"epoch: {epoch}"]
         if track_accuracy:
             parts.append(f"train acc: {train['acc'] * 100:.1f}% | val acc: {val['acc'] * 100:.1f}%")
         parts.append(f"train loss: {train['loss']:.4f} | val loss: {val['loss']:.4f}")
+        parts += [f"val {k}: {v:.3f}" for k, v in val.items() if k not in ("loss", "acc")]
         if improved and save:
             parts.append("saved")
         print(" | ".join(parts))

api/src/models/yolo.py

@@ -32,14 +32,12 @@ class SkipBlock(nn.Module):
         )
 
     def forward(self, x):
-        # F.dropout defaults to training=True, so this stays active in eval
-        # mode too. Kept as-is because the saved models were trained this way.
-        return F.dropout(F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3)
+        return F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True)
 
 
 class Yolo_model(nn.Module):
-    """Output: [batch, grid, grid, boxes*5 + labels] with per cell
-    (x, y, w, h, confidence, classes...).
+    """Output: [batch, grid x (column), grid y (row), boxes*5 + labels] with
+    per cell (x, y, w, h, confidence, classes...).
 
     size_activation decides how w/h are produced:
       "exp"     - exp(x), used by the deployed build_models/face_detection_yolo
@@ -53,6 +51,7 @@ class Yolo_model(nn.Module):
         self.size_activation = size_activation
 
         self.model = nn.Sequential(
+            #nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1, stride=2),
             nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
             nn.BatchNorm2d(c_hidden),
             nn.LeakyReLU(inplace=True),
@@ -62,27 +61,37 @@ class Yolo_model(nn.Module):
             SkipBlock(c_in=c_hidden, c_out=c_hidden),
             SkipBlock(c_in=c_hidden, c_out=c_hidden),
 
-            SkipBlock(c_in=c_hidden, c_out=c_hidden * 2),
+            #nn.Conv2d(c_hidden, c_hidden * 2, kernel_size=5, padding=2, stride=4),
+            nn.Conv2d(c_hidden, c_hidden * 2, kernel_size=3, padding=1, stride=2),
+            nn.LeakyReLU(inplace=True),
+
+            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
             SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
             SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
             SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 2),
 
-            SkipBlock(c_in=c_hidden * 2, c_out=c_hidden * 4),
+            nn.Conv2d(c_hidden * 2, c_hidden * 4, kernel_size=3, padding=1, stride=2),
+            nn.LeakyReLU(inplace=True),
+
+            SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
             SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
             SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
             SkipBlock(c_in=c_hidden * 4, c_out=c_hidden * 4),
 
-            nn.Conv2d(c_hidden * 4, c_hidden * 8, kernel_size=3, padding=1),
+            nn.Conv2d(c_hidden * 4, c_hidden * 8, kernel_size=3, padding=1, stride=2),
+            nn.LeakyReLU(inplace=True),
+
+            nn.Conv2d(c_hidden * 8, c_hidden * 8, kernel_size=3, padding=1),
             nn.BatchNorm2d(c_hidden * 8),
             nn.LeakyReLU(inplace=True),
-            nn.Dropout(0.3),
 
             nn.AdaptiveAvgPool2d((grid, grid)),
             nn.Conv2d(c_hidden * 8, boxes * 5 + labels, kernel_size=1)
         )
 
     def forward(self, x):
-        x = self.model(x).permute(0, 2, 3, 1)
+        # [batch, C, rows, cols] -> [batch, cols, rows, C]: cell [x, y] like the targets
+        x = self.model(x).permute(0, 3, 2, 1)
         if self.size_activation == "sigmoid":
             return F.sigmoid(x)
         center = F.sigmoid(x[..., :2])

api/src/production/settings.py

@@ -24,11 +24,11 @@ BIRD_MODEL_FILE = "bird_cnn"
 BIRD_IMAGE_SIZE = 64
 BIRD_C_HIDDEN = 16
 
-FACE_MODEL_FILE = "face_detection_yolo"
+FACE_MODEL_FILE = "face_detection_yolo_2"
 FACE_IMAGE_SIZE = 64
 FACE_GRID = 6
-FACE_C_HIDDEN = 16
-FACE_SIZE_ACTIVATION = "exp"   # must match how the deployed weights were trained
+FACE_C_HIDDEN = 8
+FACE_SIZE_ACTIVATION = "sigmoid"   # must match how the deployed weights were trained
 FACE_THRESHOLD = 0.9
 
 # ---- Resource limits (the server is weak, keep everything small and bounded) ----

api/src/yolo/config.py

@@ -15,6 +15,11 @@ class Config:
     image_size: int = 64
     grid: int = 6
 
+    # convert-wider: WIDER FACE -> COCO folders <wider_out_dir>/train and /test
+    wider_dir: str = ""               # contains WIDER_train/, WIDER_val/, wider_face_split/
+    wider_out_dir: str = "data/wider_face"
+    wider_min_face: float = 2.0       # drop smaller faces (px after resizing to image_size)
+
     # Model
     c_hidden: int = 16
     boxes: int = 1
@@ -24,8 +29,10 @@ class Config:
     # Training
     epochs: int = 500
     batch_size: int = 64
-    learning_rate: float = 1e-3
-    weight_decay: float = 1e-4
+    learning_rate: float = 1e-3       # peak AdamW learning rate
+    weight_decay: float = 0.05        # AdamW, conv weights only (not biases / BatchNorm)
+    warmup_epochs: int = 3            # linear warmup from learning_rate / 10
+    min_learning_rate: float = 1e-5   # cosine decay ends here at the last epoch
     train_workers: int = 2
     val_workers: int = 0
     device: str = "auto"              # "auto", "cpu", "cuda:0", ...

api/src/yolo/loss.py

@@ -13,40 +13,39 @@ class YoloLoss(nn.Module):
         target_boxes = targets[..., :4]
         target_conf = targets[..., 4]
         target_classes = targets[..., 5:]
+        N = predictions.shape[0]
 
         obj_mask = targets[..., 4] == 1
         noobj_mask = targets[..., 4] == 0
 
 
-        if obj_mask.any():
-            box_loss = F.mse_loss(
-                pred_boxes[obj_mask],
-                target_boxes[obj_mask],
-                reduction="mean"
-            )
-        else:
-            box_loss = torch.tensor(0.0, device=predictions.device)
+        box_loss = F.mse_loss(
+            pred_boxes[obj_mask],
+            target_boxes[obj_mask],
+            reduction="sum"
+        ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
 
         obj_loss = F.mse_loss(
             pred_conf[obj_mask],
             target_conf[obj_mask],
-            reduction="mean"
+            reduction="sum"
         ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
         noobj_loss = F.mse_loss(
             pred_conf[noobj_mask],
             target_conf[noobj_mask],
-            reduction="mean"
+            reduction="sum"
         ) if noobj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
 
-        class_loss = F.binary_cross_entropy_with_logits(
+        class_loss = F.binary_cross_entropy(
             pred_classes[obj_mask],
-            target_classes[obj_mask]
+            target_classes[obj_mask],
+            reduction="sum"
         ) if obj_mask.any() else torch.tensor(0.0, device=predictions.device)
 
         
-        total_loss = box_loss + obj_loss + noobj_loss + class_loss
+        total_loss = (box_loss * 5 + obj_loss + noobj_loss * 0.5 + class_loss) / N
 
         return total_loss

api/src/yolo/metrics.py

@@ -0,0 +1,51 @@
+"""Detection metrics over the whole validation set, logged next to the loss.
+Per grid cell, like the loss: a cell is a detection when its confidence
+passes `conf_threshold`, and a correct one when the cell holds a face and
+the box overlaps the true box by at least `iou_threshold`."""
+
+import torch
+
+
+def _cell_boxes(t):
+    """[N, S, S, >=4] cell outputs (x, y relative to the cell, w, h relative to
+    the image) -> [N, S, S, 4] boxes as (xmin, ymin, xmax, ymax) in image units."""
+    S = t.shape[1]
+    i = torch.arange(S, device=t.device).view(S, 1)
+    j = torch.arange(S, device=t.device).view(1, S)
+    cx = (i + t[..., 0]) / S
+    cy = (j + t[..., 1]) / S
+    w, h = t[..., 2], t[..., 3]
+    return torch.stack([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], dim=-1)
+
+
+def _iou(a, b):
+    lt = torch.max(a[..., :2], b[..., :2])
+    rb = torch.min(a[..., 2:], b[..., 2:])
+    inter = (rb - lt).clamp(min=0).prod(-1)
+    area_a = (a[..., 2:] - a[..., :2]).clamp(min=0).prod(-1)
+    area_b = (b[..., 2:] - b[..., :2]).clamp(min=0).prod(-1)
+    return inter / (area_a + area_b - inter).clamp(min=1e-9)
+
+
+def _average_precision(scores, labels):
+    """How well the confidence ranks cells with a face above empty cells (1 = perfect)."""
+    order = scores.argsort(descending=True)
+    hits = labels[order].float()
+    precision = hits.cumsum(0) / torch.arange(1, len(hits) + 1, device=hits.device)
+    return ((precision * hits).sum() / hits.sum().clamp(min=1)).item()
+
+
+def detection_metrics(preds, targets, conf_threshold=0.5, iou_threshold=0.5):
+    obj = targets[..., 4] == 1
+    conf = preds[..., 4]
+    ious = _iou(_cell_boxes(preds), _cell_boxes(targets))
+
+    detected = conf > conf_threshold
+    correct = (detected & obj & (ious > iou_threshold)).sum().item()
+    n_obj = obj.sum().item()
+    return {
+        "conf_ap": _average_precision(conf.flatten(), obj.flatten()),
+        "box_iou": ious[obj].mean().item() if n_obj else 0.0,
+        "precision": correct / max(detected.sum().item(), 1),
+        "recall": correct / max(n_obj, 1),
+    }

api/src/yolo/tasks.py

@@ -10,12 +10,14 @@ from torchvision.utils import draw_bounding_boxes
 
 from src.common.checkpoints import checkpoint_path, load_checkpoint
 from src.common.device import get_device
-from src.common.training import benchmark_num_workers, fit
+from src.common.training import adamw, benchmark_num_workers, fit, warmup_cosine
 from src.common.visualize import show_image
 from src.models.yolo import Yolo_model, convert_prediction
 from src.yolo.config import Config
 from src.yolo.dataset import YoloDataset
 from src.yolo.loss import YoloLoss
+from src.yolo.metrics import detection_metrics
+from src.yolo.wider import convert_wider
 
 
 def build_dataset(cfg: Config, folder: str, augment: bool):
@@ -50,7 +52,7 @@ def draw_prediction(label, image, threshold):
 
 
 def train(cfg: Config):
-    """Train the detector; saves the model with the lowest validation loss."""
+    """Train the detector; saves the model with the highest validation conf_ap."""
     device = get_device(cfg.device)
     train_loader = DataLoader(build_dataset(cfg, cfg.train_dir, cfg.augment_train),
                               batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.train_workers)
@@ -58,11 +60,12 @@ def train(cfg: Config):
                             batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.val_workers)
 
     model = build_model(cfg).to(device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
+    optimizer = adamw(model, cfg.learning_rate, cfg.weight_decay)
+    scheduler = warmup_cosine(optimizer, cfg.epochs, cfg.warmup_epochs, cfg.min_learning_rate)
 
     fit(model, YoloLoss(), optimizer, train_loader, val_loader, device,
         epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.model_name, save=cfg.save,
-        monitor="loss", config=cfg)
+        monitor="conf_ap", val_metrics=detection_metrics, config=cfg, scheduler=scheduler)
 
 
 def view_data(cfg: Config):
@@ -146,4 +149,5 @@ TASKS = {
     "sample": sample,
     "benchmark-workers": benchmark_workers,
     "webcam": webcam,
+    "convert-wider": convert_wider,
 }

api/src/yolo/wider.py

@@ -0,0 +1,135 @@
+"""Convert the WIDER FACE dataset to the COCO folders YoloDataset reads.
+
+WIDER_train becomes <wider_out_dir>/train and WIDER_val becomes
+<wider_out_dir>/test (WIDER's own test split has no public boxes). Each folder
+gets its images (copied, flat) and <annotation_file> with one category,
+1 = face, like faces_scenery.
+
+Most WIDER faces are tiny faces in crowds, which the model cannot use, so:
+- faces smaller than --wider-min-face px after resizing to --image-size are dropped,
+- images where two faces fall in the same --grid cell are skipped. YoloDataset
+  keeps one face per cell, so the other face would be trained as background.
+"""
+
+import json
+import shutil
+from pathlib import Path
+
+import cv2
+from tqdm import tqdm
+
+from src.common.boxes import grid_cell
+from src.yolo.config import Config
+
+# WIDER split -> output folder (train_dir / val_dir)
+SPLITS = {"train": "train", "val": "test"}
+
+
+def find(root: Path, pattern: str) -> Path:
+    """First match of pattern below root, so nested unzips (WIDER_train/WIDER_train/...) work too."""
+    matches = sorted(root.rglob(pattern))
+    if not matches:
+        raise FileNotFoundError(f"{pattern} not found under {root}")
+    return matches[0]
+
+
+def parse_ground_truth(path: Path):
+    """Yield (image path, faces), one face per row as
+    [x, y, w, h, blur, expression, illumination, invalid, occlusion, pose]."""
+    lines = [line.strip() for line in path.read_text().splitlines() if line.strip()]
+    i = 0
+    while i < len(lines):
+        name, count = lines[i], int(lines[i + 1])
+        faces = [[int(v) for v in line.split()] for line in lines[i + 2 : i + 2 + count]]
+        i += 2 + max(count, 1)  # an image without faces still has one "0 0 0 ..." row
+        yield name, faces
+
+
+def usable_boxes(faces, img_w, img_h, image_size, min_face):
+    """WIDER faces -> COCO [xmin, ymin, w, h] boxes in whole pixels, clipped to
+    the image and at least min_face px wide and high once resized to image_size."""
+    boxes = []
+    for x, y, w, h, _blur, _expression, _illumination, invalid, _occlusion, _pose in faces:
+        x0, y0 = max(x, 0), max(y, 0)
+        w, h = min(x + w, img_w) - x0, min(y + h, img_h) - y0
+        if invalid or w < 1 or h < 1:
+            continue
+        if w * image_size / img_w < min_face or h * image_size / img_h < min_face:
+            continue
+        boxes.append([x0, y0, w, h])
+    return boxes
+
+
+def shares_grid_cell(boxes, img_w, img_h, image_size, grid, eps=1e-3):
+    """True when two box centres can land in the same grid cell after resizing
+    to image_size. A centre within eps px of a cell border counts for both
+    cells, since Albumentations' float32 math can round it to either side."""
+    sx, sy = image_size / img_w, image_size / img_h
+    seen = set()
+    for x, y, w, h in boxes:
+        cx, cy = (x + w / 2) * sx, (y + h / 2) * sy
+        cells = {grid_cell(cx + dx, cy + dy, image_size, image_size, grid)
+                 for dx in (-eps, eps) for dy in (-eps, eps)}
+        if cells & seen:
+            return True
+        seen |= cells
+    return False
+
+
+def convert_split(images_dir: Path, gt_file: Path, out_dir: Path, cfg: Config):
+    out_dir.mkdir(parents=True)
+    images, annotations = [], []
+    skipped = {"unreadable": 0, "faces sharing a cell": 0}
+
+    for name, faces in tqdm(list(parse_ground_truth(gt_file)), desc=out_dir.name):
+        source = images_dir / name
+        # read with OpenCV like YoloDataset does, so every kept image is loadable
+        image = cv2.imread(str(source))
+        if image is None:
+            skipped["unreadable"] += 1
+            continue
+        img_h, img_w = image.shape[:2]
+
+        boxes = usable_boxes(faces, img_w, img_h, cfg.image_size, cfg.wider_min_face)
+        if shares_grid_cell(boxes, img_w, img_h, cfg.image_size, cfg.grid):
+            skipped["faces sharing a cell"] += 1
+            continue
+
+        target = out_dir / source.name
+        if target.exists():
+            raise FileExistsError(f"two WIDER images are named {source.name}")
+        shutil.copy2(source, target)
+
+        image_id = len(images) + 1
+        images.append({"id": image_id, "file_name": source.name, "width": img_w, "height": img_h})
+        for box in boxes:
+            annotations.append({"id": len(annotations) + 1, "image_id": image_id, "category_id": 1,
+                                "bbox": box, "area": box[2] * box[3], "iscrowd": 0})
+
+    coco = {"images": images, "annotations": annotations,
+            "categories": [{"id": 1, "name": "face", "supercategory": ""}]}
+    (out_dir / cfg.annotation_file).write_text(json.dumps(coco))
+
+    empty = len(images) - len({a["image_id"] for a in annotations})
+    print(f"{out_dir}: {len(images)} images ({empty} without faces), {len(annotations)} faces. "
+          f"Skipped: " + ", ".join(f"{n} {reason}" for reason, n in skipped.items()))
+
+
+def convert_wider(cfg: Config):
+    """Convert WIDER FACE (--wider-dir) to COCO train/ and test/ folders in --wider-out-dir."""
+    if not cfg.wider_dir:
+        raise SystemExit("Set --wider-dir to the folder with WIDER_train/, WIDER_val/ and wider_face_split/")
+    root, out = Path(cfg.wider_dir), Path(cfg.wider_out_dir)
+
+    sources = {split: (find(root, f"WIDER_{split}/images"), find(root, f"wider_face_{split}_bbx_gt.txt"))
+               for split in SPLITS}
+    for folder in SPLITS.values():
+        if (out / folder).exists():
+            raise SystemExit(f"{out / folder} already exists, delete it or pick another --wider-out-dir")
+
+    for split, folder in SPLITS.items():
+        images_dir, gt_file = sources[split]
+        convert_split(images_dir, gt_file, out / folder, cfg)
+
+    print(f"Train with: --train-dir {out / 'train'} --val-dir {out / 'test'} "
+          f"--image-size {cfg.image_size} --grid {cfg.grid} --labels 1")

api/train_test.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+
+for ((c = 1; c <= 32; c *= 2)); do
+    python -m src.main yolo.train --train-dir ./data/wider_faces_coco/train/ --val-dir ./data/wider_faces_coco/test/ \
+        --no-augment-val --c-hidden "$c" --batch-size 64 --epochs 200 --run-name "bn-wider-faces-c$c" \
+        --no-save --train-workers 4 --image-size 256 --grid 8 \
+        --learning-rate 1e-3 --weight-decay 0.05 --warmup-epochs 3 --min-learning-rate 1e-5
+done
+
+
+sudo -n shutdown -h +2