Compare commits
2
Commits
3faf0c28ca
...
0715c66a64
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0715c66a64 | ||
|
|
c36f34dea9 |
@@ -1,6 +1,7 @@
|
|||||||
api/data/
|
api/data/
|
||||||
api/data.*
|
api/data.*
|
||||||
api/saved_models/
|
api/saved_models/
|
||||||
|
api/runs/
|
||||||
api/testimages/
|
api/testimages/
|
||||||
**/__pycache__/
|
**/__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"
|
model_name: str = "bird_cnn"
|
||||||
save: bool = True
|
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
|
# predict task
|
||||||
image: str = "test1.jpg"
|
image: str = "test1.jpg"
|
||||||
checkpoint: str = "" # empty = <save_dir>/<model_name>/<model_name>
|
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,
|
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,
|
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):
|
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
|
Each epoch trains, validates, prints the metrics and saves the model when the
|
||||||
monitored metric improves. `prepare_batch` turns a loader batch into
|
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 math
|
||||||
import time
|
import time
|
||||||
@@ -12,6 +15,7 @@ from torch.utils.data import DataLoader
|
|||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from src.common.checkpoints import save_checkpoint
|
from src.common.checkpoints import save_checkpoint
|
||||||
|
from src.common.run_log import RunLog
|
||||||
|
|
||||||
|
|
||||||
def _default_prepare(batch, device):
|
def _default_prepare(batch, device):
|
||||||
@@ -64,10 +68,34 @@ def fit(
|
|||||||
monitor: str = "loss", # "loss" (lower is better) or "acc" (higher is better)
|
monitor: str = "loss", # "loss" (lower is better) or "acc" (higher is better)
|
||||||
track_accuracy: bool = False, # for classifiers
|
track_accuracy: bool = False, # for classifiers
|
||||||
prepare_batch=_default_prepare,
|
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
|
best = math.inf if monitor == "loss" else -math.inf
|
||||||
|
|
||||||
for epoch in range(1, epochs + 1):
|
for epoch in range(1, epochs + 1):
|
||||||
|
start = time.perf_counter()
|
||||||
train = _run_epoch(model, loss_fn, train_loader, device, prepare_batch, track_accuracy,
|
train = _run_epoch(model, loss_fn, train_loader, device, prepare_batch, track_accuracy,
|
||||||
optimizer=optimizer, desc=f"Train {epoch}")
|
optimizer=optimizer, desc=f"Train {epoch}")
|
||||||
val = _run_epoch(model, loss_fn, val_loader, device, prepare_batch, track_accuracy,
|
val = _run_epoch(model, loss_fn, val_loader, device, prepare_batch, track_accuracy,
|
||||||
@@ -78,6 +106,11 @@ def fit(
|
|||||||
best = val[monitor]
|
best = val[monitor]
|
||||||
if save:
|
if save:
|
||||||
save_checkpoint(model, save_dir, model_name)
|
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.
|
# Loss is the summed batch loss per sample, shown x1000.
|
||||||
parts = [f"epoch: {epoch}"]
|
parts = [f"epoch: {epoch}"]
|
||||||
|
|||||||
@@ -14,20 +14,20 @@ class SkipBlock(nn.Module):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self.conv = nn.Sequential(
|
self.conv = nn.Sequential(
|
||||||
nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size // 2),
|
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.LeakyReLU(inplace=True),
|
||||||
|
|
||||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2),
|
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.LeakyReLU(inplace=True),
|
||||||
|
|
||||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size // 2),
|
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.LeakyReLU(inplace=True)
|
||||||
)
|
)
|
||||||
self.conv_skip = nn.Sequential(
|
self.conv_skip = nn.Sequential(
|
||||||
nn.Conv2d(c_in, c_out, 1),
|
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)
|
nn.LeakyReLU(inplace=True)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -54,7 +54,7 @@ class Yolo_model(nn.Module):
|
|||||||
|
|
||||||
self.model = nn.Sequential(
|
self.model = nn.Sequential(
|
||||||
nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
|
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),
|
nn.LeakyReLU(inplace=True),
|
||||||
|
|
||||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
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),
|
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),
|
||||||
nn.GroupNorm(num_groups=c_hidden // 2, num_channels=c_hidden * 8),
|
nn.BatchNorm2d(c_hidden * 8),
|
||||||
nn.LeakyReLU(inplace=True),
|
nn.LeakyReLU(inplace=True),
|
||||||
nn.Dropout(0.3),
|
nn.Dropout(0.3),
|
||||||
|
|
||||||
|
|||||||
@@ -40,6 +40,10 @@ class Config:
|
|||||||
football_model_name: str = "object_detection"
|
football_model_name: str = "object_detection"
|
||||||
save: bool = True
|
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
|
# detect tasks: proposals smaller than min_size or larger than max_size are ignored
|
||||||
image: str = "" # detect-faces: image to run on
|
image: str = "" # detect-faces: image to run on
|
||||||
min_size: int = 5
|
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,
|
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,
|
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):
|
def detect_faces(cfg: Config):
|
||||||
@@ -103,7 +103,7 @@ def train_football(cfg: Config):
|
|||||||
|
|
||||||
fit(model, nn.CrossEntropyLoss(), optimizer, train_loader, val_loader, device,
|
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,
|
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):
|
def detect_football(cfg: Config):
|
||||||
|
|||||||
@@ -35,6 +35,10 @@ class Config:
|
|||||||
model_name: str = "face_detection_yolo"
|
model_name: str = "face_detection_yolo"
|
||||||
save: bool = True
|
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
|
# sample / webcam tasks
|
||||||
checkpoint: str = "" # empty = <save_dir>/<model_name>/<model_name>
|
checkpoint: str = "" # empty = <save_dir>/<model_name>/<model_name>
|
||||||
threshold: float = 0.95 # confidence needed to draw a box
|
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,
|
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,
|
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):
|
def view_data(cfg: Config):
|
||||||
|
|||||||
Reference in New Issue
Block a user