Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0715c66a64 | ||
|
|
c36f34dea9 |
@@ -1,6 +1,7 @@
|
||||
api/data/
|
||||
api/data.*
|
||||
api/saved_models/
|
||||
api/runs/
|
||||
api/testimages/
|
||||
**/__pycache__/
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+325
@@ -0,0 +1,325 @@
|
||||
# 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.
|
||||
@@ -28,6 +28,10 @@ class Config:
|
||||
model_name: str = "bird_cnn"
|
||||
save: bool = True
|
||||
|
||||
# Run log (src/common/run_log.py): one folder per training run
|
||||
runs_dir: str = "runs" # empty = no run log
|
||||
run_name: str = "" # added to the run folder name, e.g. "bn-lr3e-4"
|
||||
|
||||
# predict task
|
||||
image: str = "test1.jpg"
|
||||
checkpoint: str = "" # empty = <save_dir>/<model_name>/<model_name>
|
||||
|
||||
@@ -56,7 +56,7 @@ def train(cfg: Config):
|
||||
|
||||
fit(model, nn.CrossEntropyLoss(), optimizer, train_loader, val_loader, device,
|
||||
epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.model_name, save=cfg.save,
|
||||
monitor="acc", track_accuracy=True)
|
||||
monitor="acc", track_accuracy=True, config=cfg)
|
||||
|
||||
|
||||
def view_data(cfg: Config):
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
"""Records every training run in its own folder, so runs with different
|
||||
settings or architectures can be compared later.
|
||||
|
||||
<runs_dir>/<model_name>/<YYYYmmdd-HHMMSS>[_<run_name>]/
|
||||
run.json everything about the run: config, command, model
|
||||
(class, parameter counts, architecture hash),
|
||||
optimizer, loss, data sizes, environment, git state
|
||||
and the result (status, best epoch and metric,
|
||||
duration). Rewritten after every epoch.
|
||||
metrics.csv one row per epoch
|
||||
architecture.txt the model's module tree, str(model)
|
||||
changes.diff uncommitted code changes (only if there are any)
|
||||
best.pt state_dict of the best epoch so far
|
||||
|
||||
result.status in run.json is "running" during training, then "completed",
|
||||
"failed" or "interrupted". A run whose process was killed stays "running".
|
||||
"""
|
||||
|
||||
import csv
|
||||
import dataclasses
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
REPO_DIR = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _git(*args):
|
||||
try:
|
||||
return subprocess.run(["git", *args], cwd=REPO_DIR, capture_output=True, text=True,
|
||||
timeout=10, check=True).stdout
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return None
|
||||
|
||||
|
||||
def _jsonable(value):
|
||||
if value is None or isinstance(value, (str, int, float, bool)):
|
||||
return value
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [_jsonable(v) for v in value]
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _jsonable(v) for k, v in value.items()}
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.tolist() if value.numel() <= 16 else f"tensor{tuple(value.shape)}"
|
||||
return str(value)
|
||||
|
||||
|
||||
def _module_settings(module):
|
||||
"""Public attributes and small buffers, e.g. a loss's reduction and class weights."""
|
||||
settings = {k: v for k, v in vars(module).items() if not k.startswith("_") and k != "training"}
|
||||
settings.update(dict(module.named_buffers()))
|
||||
return _jsonable(settings)
|
||||
|
||||
|
||||
def _dataset_size(loader):
|
||||
try:
|
||||
return len(loader.dataset)
|
||||
except TypeError:
|
||||
return None
|
||||
|
||||
|
||||
class RunLog:
|
||||
def __init__(self, runs_dir, model_name, config, model, optimizer, loss_fn,
|
||||
train_loader, val_loader, device, monitor):
|
||||
self.monitor = monitor
|
||||
self.started = time.time()
|
||||
self.dir = self._make_dir(Path(runs_dir) / model_name, getattr(config, "run_name", ""))
|
||||
self._csv_fields = None
|
||||
|
||||
architecture = str(model)
|
||||
status = _git("status", "--porcelain")
|
||||
diff = _git("diff", "HEAD")
|
||||
gpu = torch.cuda.get_device_name(device) if torch.device(device).type == "cuda" else None
|
||||
|
||||
self.info = {
|
||||
"id": f"{model_name}/{self.dir.name}",
|
||||
"model_name": model_name,
|
||||
"started": datetime.fromtimestamp(self.started).isoformat(timespec="seconds"),
|
||||
"command": " ".join(sys.argv),
|
||||
"config": _jsonable(dataclasses.asdict(config)),
|
||||
"model": {
|
||||
"class": type(model).__name__,
|
||||
"params": sum(p.numel() for p in model.parameters()),
|
||||
"trainable_params": sum(p.numel() for p in model.parameters() if p.requires_grad),
|
||||
"architecture_hash": hashlib.sha256(architecture.encode()).hexdigest()[:12],
|
||||
},
|
||||
"optimizer": {
|
||||
"class": type(optimizer).__name__,
|
||||
**_jsonable({k: v for k, v in optimizer.param_groups[0].items() if k != "params"}),
|
||||
},
|
||||
"loss": {"class": type(loss_fn).__name__, **_module_settings(loss_fn)},
|
||||
"data": {
|
||||
"train_size": _dataset_size(train_loader),
|
||||
"val_size": _dataset_size(val_loader),
|
||||
"batch_size": train_loader.batch_size,
|
||||
"train_workers": train_loader.num_workers,
|
||||
},
|
||||
"environment": {
|
||||
"host": socket.gethostname(),
|
||||
"platform": platform.platform(),
|
||||
"python": platform.python_version(),
|
||||
"torch": torch.__version__,
|
||||
"cuda": torch.version.cuda,
|
||||
"device": str(device),
|
||||
"gpu": gpu,
|
||||
"seed": torch.initial_seed(),
|
||||
},
|
||||
"git": {
|
||||
"commit": (_git("rev-parse", "HEAD") or "").strip() or None,
|
||||
"branch": (_git("rev-parse", "--abbrev-ref", "HEAD") or "").strip() or None,
|
||||
"dirty": bool(status.strip()) if status is not None else None,
|
||||
"changed_files": [line[3:] for line in status.splitlines()] if status else [],
|
||||
},
|
||||
"result": {
|
||||
"status": "running",
|
||||
"monitor": monitor,
|
||||
"epochs_done": 0,
|
||||
"best_epoch": None,
|
||||
f"best_val_{monitor}": None,
|
||||
"last": None,
|
||||
"duration_s": 0,
|
||||
"finished": None,
|
||||
"error": None,
|
||||
},
|
||||
}
|
||||
|
||||
(self.dir / "architecture.txt").write_text(architecture + "\n")
|
||||
if diff:
|
||||
(self.dir / "changes.diff").write_text(diff)
|
||||
self._write_info()
|
||||
print(f"Logging run to {self.dir}")
|
||||
|
||||
@staticmethod
|
||||
def _make_dir(parent, run_name):
|
||||
name = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
if run_name:
|
||||
name += "_" + re.sub(r"[^\w.-]+", "-", run_name)
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
for i in range(1, 100):
|
||||
path = parent / (name if i == 1 else f"{name}-{i}")
|
||||
try:
|
||||
path.mkdir()
|
||||
return path
|
||||
except FileExistsError:
|
||||
continue
|
||||
raise FileExistsError(f"no free run folder for {parent / name}")
|
||||
|
||||
def _write_info(self):
|
||||
# Write-then-rename so a reader never sees a half written file.
|
||||
tmp = self.dir / "run.json.tmp"
|
||||
tmp.write_text(json.dumps(self.info, indent=2) + "\n")
|
||||
os.replace(tmp, self.dir / "run.json")
|
||||
|
||||
def log_epoch(self, epoch, train, val, lr, seconds, improved):
|
||||
row = {"epoch": epoch,
|
||||
**{f"train_{k}": v for k, v in train.items()},
|
||||
**{f"val_{k}": v for k, v in val.items()},
|
||||
"lr": lr, "seconds": round(seconds, 2), "improved": improved}
|
||||
path = self.dir / "metrics.csv"
|
||||
with path.open("a", newline="") as f:
|
||||
if self._csv_fields is None:
|
||||
self._csv_fields = list(row)
|
||||
writer = csv.DictWriter(f, fieldnames=self._csv_fields)
|
||||
writer.writeheader()
|
||||
else:
|
||||
writer = csv.DictWriter(f, fieldnames=self._csv_fields)
|
||||
writer.writerow(row)
|
||||
|
||||
result = self.info["result"]
|
||||
result["epochs_done"] = epoch
|
||||
result["last"] = {k: v for k, v in row.items() if k != "improved"}
|
||||
if improved:
|
||||
result["best_epoch"] = epoch
|
||||
result[f"best_val_{self.monitor}"] = val[self.monitor]
|
||||
result["duration_s"] = round(time.time() - self.started, 1)
|
||||
self._write_info()
|
||||
|
||||
def save_weights(self, model):
|
||||
torch.save(model.state_dict(), self.dir / "best.pt")
|
||||
|
||||
def finish(self, status, error=None):
|
||||
result = self.info["result"]
|
||||
result["status"] = status
|
||||
result["error"] = error
|
||||
result["duration_s"] = round(time.time() - self.started, 1)
|
||||
result["finished"] = datetime.now().isoformat(timespec="seconds")
|
||||
self._write_info()
|
||||
@@ -2,7 +2,10 @@
|
||||
|
||||
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."""
|
||||
(inputs, targets); the default just moves both to the device.
|
||||
|
||||
Passing the task's `config` also records the run (settings, architecture,
|
||||
metrics per epoch, best weights) under config.runs_dir, see run_log.py."""
|
||||
|
||||
import math
|
||||
import time
|
||||
@@ -12,6 +15,7 @@ from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from src.common.checkpoints import save_checkpoint
|
||||
from src.common.run_log import RunLog
|
||||
|
||||
|
||||
def _default_prepare(batch, device):
|
||||
@@ -64,10 +68,34 @@ def fit(
|
||||
monitor: str = "loss", # "loss" (lower is better) or "acc" (higher is better)
|
||||
track_accuracy: bool = False, # for classifiers
|
||||
prepare_batch=_default_prepare,
|
||||
config=None, # the task's Config; enables the run log
|
||||
):
|
||||
run = None
|
||||
if config is not None and config.runs_dir:
|
||||
run = RunLog(config.runs_dir, model_name, config, model, optimizer, loss_fn,
|
||||
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)
|
||||
except KeyboardInterrupt:
|
||||
if run:
|
||||
run.finish("interrupted")
|
||||
raise
|
||||
except Exception as e:
|
||||
if run:
|
||||
run.finish("failed", error=repr(e))
|
||||
raise
|
||||
if run:
|
||||
run.finish("completed")
|
||||
return best
|
||||
|
||||
|
||||
def _fit_epochs(model, loss_fn, optimizer, train_loader, val_loader, device, epochs,
|
||||
save_dir, model_name, save, monitor, track_accuracy, prepare_batch, run):
|
||||
best = math.inf if monitor == "loss" else -math.inf
|
||||
|
||||
for epoch in range(1, epochs + 1):
|
||||
start = time.perf_counter()
|
||||
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,
|
||||
@@ -78,6 +106,11 @@ def fit(
|
||||
best = val[monitor]
|
||||
if save:
|
||||
save_checkpoint(model, save_dir, model_name)
|
||||
if run:
|
||||
run.save_weights(model)
|
||||
if run:
|
||||
run.log_epoch(epoch, train, val, lr=optimizer.param_groups[0]["lr"],
|
||||
seconds=time.perf_counter() - start, improved=improved)
|
||||
|
||||
# Loss is the summed batch loss per sample, shown x1000.
|
||||
parts = [f"epoch: {epoch}"]
|
||||
|
||||
@@ -14,20 +14,20 @@ class SkipBlock(nn.Module):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size // 2),
|
||||
nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2),
|
||||
nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2),
|
||||
nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.LeakyReLU(inplace=True)
|
||||
)
|
||||
self.conv_skip = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1),
|
||||
nn.GroupNorm(num_groups=c_out // 8, num_channels=c_out),
|
||||
nn.BatchNorm2d(c_out),
|
||||
nn.LeakyReLU(inplace=True)
|
||||
)
|
||||
|
||||
@@ -54,7 +54,7 @@ class Yolo_model(nn.Module):
|
||||
|
||||
self.model = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(num_groups=c_hidden // 8, num_channels=c_hidden),
|
||||
nn.BatchNorm2d(c_hidden),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
@@ -73,7 +73,7 @@ class Yolo_model(nn.Module):
|
||||
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.GroupNorm(num_groups=c_hidden // 2, num_channels=c_hidden * 8),
|
||||
nn.BatchNorm2d(c_hidden * 8),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
nn.Dropout(0.3),
|
||||
|
||||
|
||||
@@ -40,6 +40,10 @@ class Config:
|
||||
football_model_name: str = "object_detection"
|
||||
save: bool = True
|
||||
|
||||
# Run log (src/common/run_log.py): one folder per training run
|
||||
runs_dir: str = "runs" # empty = no run log
|
||||
run_name: str = "" # added to the run folder name, e.g. "bn-lr3e-4"
|
||||
|
||||
# detect tasks: proposals smaller than min_size or larger than max_size are ignored
|
||||
image: str = "" # detect-faces: image to run on
|
||||
min_size: int = 5
|
||||
|
||||
@@ -71,7 +71,7 @@ def train_faces(cfg: Config):
|
||||
|
||||
fit(model, nn.CrossEntropyLoss(weight=weights), optimizer, train_loader, val_loader, device,
|
||||
epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.faces_model_name, save=cfg.save,
|
||||
monitor="loss", track_accuracy=True)
|
||||
monitor="loss", track_accuracy=True, config=cfg)
|
||||
|
||||
|
||||
def detect_faces(cfg: Config):
|
||||
@@ -103,7 +103,7 @@ def train_football(cfg: Config):
|
||||
|
||||
fit(model, nn.CrossEntropyLoss(), optimizer, train_loader, val_loader, device,
|
||||
epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.football_model_name, save=cfg.save,
|
||||
monitor="loss", track_accuracy=True, prepare_batch=prepare_batch)
|
||||
monitor="loss", track_accuracy=True, prepare_batch=prepare_batch, config=cfg)
|
||||
|
||||
|
||||
def detect_football(cfg: Config):
|
||||
|
||||
@@ -35,6 +35,10 @@ class Config:
|
||||
model_name: str = "face_detection_yolo"
|
||||
save: bool = True
|
||||
|
||||
# Run log (src/common/run_log.py): one folder per training run
|
||||
runs_dir: str = "runs" # empty = no run log
|
||||
run_name: str = "" # added to the run folder name, e.g. "bn-lr3e-4"
|
||||
|
||||
# sample / webcam tasks
|
||||
checkpoint: str = "" # empty = <save_dir>/<model_name>/<model_name>
|
||||
threshold: float = 0.95 # confidence needed to draw a box
|
||||
|
||||
@@ -62,7 +62,7 @@ def train(cfg: Config):
|
||||
|
||||
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")
|
||||
monitor="loss", config=cfg)
|
||||
|
||||
|
||||
def view_data(cfg: Config):
|
||||
|
||||
Reference in New Issue
Block a user