Add training support: fine-tuning, from-scratch, v3/v4/v7 families

[IanButterworth]

Note

This PR was written by Claude (Claude Code), from training support requested and supervised by @IanButterworth.

Adds working training support to ObjectDetector, replacing the README's "Training using ObjectDetector is currently unproven/untested". Supports fine-tuning pretrained weights and from-scratch training, for the v3, v4, and v7 model families.

What's included

• train!(model, data; ...) — full training loop: batching/shuffling, Adam (or any Optimisers rule), LR warmup (darknet burn-in style), optional horizontal-flip augmentation, per-epoch loss reporting, and checkpointing to darknet-format .weights.
• Differentiable forward path — the buffered inference forward mutates preallocated arrays and isn't AD-able, so the layer closures became callable structs (RouteLayer, ShortcutAdd, CatLayer, …) that keep the exact fast inference path but also record their source chain indices, letting a pure non-mutating forward resolve skips functionally. Verified numerically identical to the buffered forward.
• Loss — CIoU box loss + BCE objectness/class with single-best-anchor assignment (modern simplified, not darknet's exact loss). The box loss adds a raw-space wh regression term sharing the same optimum as CIoU: CIoU alone has vanishing wh gradients once the classic exp decode over/undershoots, which let predictions run away irrecoverably (found while producing the session below). Classic heads (v3 family, v4, v4-tiny: linear/logit) and new_coords=1 heads (scaled-YOLOv4 + v7 families: logistic/probability, (t·sxy)² decode) are both supported; yolov2 [region] models error cleanly.
• From-scratch training — trainable_batchnorm=true keeps live BatchNorm layers instead of folding them into the convs (default remains folding, which is faster for inference but freezes BN statistics). Verified output-equivalent to the folded path on pretrained weights.
• Transfer learning — weights_stop_layer=N random-inits everything past a backbone split (15 ≙ darknet's yolov3-tiny.conv.15), enabling custom class counts via the existing cfgchanges; allow_partial_weights=true loads truncated .conv.XX files.
• save_weights — darknet-format output that round-trips with the same cfg: real BN parameters in trainable_batchnorm mode, identity-BN parameters for folded models.
• Data handling — TrainSample, load_darknet_dataset / load_darknet_labels (darknet label conventions), letterbox transforms applied to boxes consistently with prepare_image.

Example session

Fine-tuning the pretrained COCO backbone to a single custom class ("white rectangle") at 160×160, on CPU. Output below is a genuine captured run.

julia> using ObjectDetector

julia> cfg, weights = YOLO.YOLO_MODELS["v3_tiny_COCO"]();

julia> yolomod = YOLO.Yolo(cfg, weights, 1;
           silent=true,
           weights_stop_layer = 15,  # keep pretrained backbone (≙ yolov3-tiny.conv.15), random-init the heads
           cfgchanges = [(:net, 1, :width, 160), (:net, 1, :height, 160),
                         (:yolo, 1, :classes, 1), (:yolo, 2, :classes, 1),
                         (:convolutional, 10, :filters, 18), (:convolutional, 13, :filters, 18)]);

julia> mkimg(yr, xr) = (img = fill(0.1f0, 160, 160, 3); img[yr, xr, :] .= 1f0; img);

julia> mkbox(yr, xr) = Float32[1; (first(xr)+last(xr))/2/160; (first(yr)+last(yr))/2/160; length(xr)/160; length(yr)/160;;];

julia> rects = [(30:60, 40:80), (100:140, 20:50), (70:110, 90:130), (20:45, 100:150)];

julia> data = [TrainSample(mkimg(yr, xr), mkbox(yr, xr)) for (yr, xr) in rects];

julia> batch = emptybatch(yolomod); batch[:, :, :, 1], _ = prepare_image(data[1].image, yolomod);

julia> yolomod(batch; detect_thresh=0.5)   # before training: no detections
10×0 Matrix{Float32}

julia> result = train!(yolomod, data; epochs=80, batchsize=4, lr=1e-3,
                       warmup_batches=5, noobj_weight=2.0);
epoch 1/80  loss: 5.0675125
epoch 2/80  loss: 4.049077
epoch 3/80  loss: 3.141153
epoch 4/80  loss: 2.1888294
epoch 5/80  loss: 2.2226012
epoch 6/80  loss: 1.7065738
epoch 7/80  loss: 1.3220385
epoch 8/80  loss: 1.131324
epoch 9/80  loss: 0.9857571
epoch 10/80  loss: 0.6409655
⋮
epoch 77/80  loss: 0.08988118
epoch 78/80  loss: 0.112671584
epoch 79/80  loss: 0.19564329
epoch 80/80  loss: 0.14985168

julia> yolomod(batch; detect_thresh=0.5)   # after: one detection, matching the truth box
10×1 Matrix{Float32}:
 0.26203835   # x1  (truth 0.25)
 0.18493523   # y1  (truth 0.1875)
 0.50191855   # x2  (truth 0.5)
 0.38029766   # y2  (truth 0.375)
 0.997701     # objectness
 0.99682426   # class confidence
 0.0
 0.99682426
 1.0          # class 1
 1.0          # batch 1

julia> save_weights(yolomod, "rectangles.weights");  # darknet-format, reloads with the same cfg

From-scratch works the same way with weights_stop_layer=0, trainable_batchnorm=true — the test suite overfits yolov3-tiny, yolov4-tiny, and yolov7-tiny from random init this way.

Inference is unchanged

• All 390 Darknet-parity tests pass unchanged; the struct layers execute the same mutating fast path the closures did.
• Forward-pass benchmarks (CPU, 416×416, min of 20+ runs) show no regression:

Model	master	this PR
v3_tiny_COCO	32.6 ms	32.6 ms
v4_tiny_COCO	44.5 ms	44.3–44.7 ms
v7_tiny_COCO	65.7 ms	65.1 ms
v3_COCO	284.0 ms	282.2 ms

Training throughput vs darknet

Benchmarked train! against darknet's own native training (train_network: forward + backward + update), driven through Darknet_jll's libdarknet — the same binary the test suite trusts for inference parity — via a new training API in Darknet.jl (training_data / train_network! ccall wrappers, to be PR'd there). Both sides trained the identical cfg (yolov3-tiny @ 416×416, batch=8, subdivisions=1) from the same pretrained weights on the identical batch content, no augmentation, CPU (macOS arm64; Julia BLAS 6 threads).

sec / batch-8 iteration	min	median	throughput
darknet train_network	28.18 s	28.96 s	0.3 img/s
ObjectDetector train! step	1.19 s	1.36 s	5.9 img/s

~21× faster here, but read with care: the jll libdarknet is a generic portable build, and darknet's real-world training performance comes from natively tuned builds (AVX on x86) or CUDA — this comparison says Julia/Flux training throughput is strongly competitive on CPU, not that darknet is inherently slow. Correctness of the darknet side was validated separately (its loss decreases and per-layer IoU climbs on toy data, and its save_weights output round-trips).

A convergence-quality comparison (same dataset trained both ways, detections evaluated identically) is prepared but not yet run.

Notes

• The loss is deliberately a modern simplified one, not a reimplementation of darknet's per-cfg loss options (iou_loss, ignore_thresh semantics, multi-anchor assignment, mosaic etc. are not implemented).
• GPU training path exists (targets/offsets move with the model) but has only been exercised on CPU here.
• Training tests run from-scratch overfitting on toy data in ~3 min on CPU (69 tests).

🤖 Generated with Claude Code