new build model

api/README.md

@@ -92,7 +92,7 @@ api/
 │   │   ├── run_log.py          records every training run to runs/
 │   │   └── visualize.py        matplotlib helpers (lazy import)
 │   ├── bird_cnn/               "bird" task group:   config.py, tasks.py
-│   ├── yolo/                   "yolo" task group:   config.py, tasks.py, dataset.py, loss.py, metrics.py
+│   ├── yolo/                   "yolo" task group:   config.py, tasks.py, dataset.py, loss.py, metrics.py, wider.py
 │   ├── r_cnn/                  "rcnn" task group:   config.py, tasks.py, dataset.py, crops.py, detect.py
 │   └── production/             "server" task group + the FastAPI app
 │       ├── app.py              FastAPI app, CORS, startup (load models, create DB)
@@ -317,14 +317,19 @@ Code: `src/yolo/` · Defaults: `src/yolo/config.py` · Model: [`Yolo_model`](#yo
 | `--augment-val` / `--no-augment-val` | **on** | train, view-data, sample | Augment validation images too (historical default, see [quirks](#known-issues-and-quirks)). |
 | `--image-size` | `64` | all | Model input size (square). |
 | `--grid` | `6` | all | Grid size S: the model predicts an S × S grid. |
+| `--wider-dir` | `""` | convert-wider | WIDER FACE folder with `WIDER_train/`, `WIDER_val/` and `wider_face_split/`. |
+| `--wider-out-dir` | `data/wider_face` | convert-wider | Output folder; gets `train/` and `test/`. |
+| `--wider-min-face` | `2.0` | convert-wider | Drop faces smaller than this (px, after resizing to `--image-size`). |
 | `--c-hidden` | `16` | all | Base channel count (c, 2c, 4c, 8c). |
 | `--boxes` | `1` | all | Boxes per cell. Only 1 is supported by the dataset, loss and decoder. |
 | `--labels` | `1` | all | Number of classes. Must equal the number of categories in the COCO file. |
 | `--size-activation` | `sigmoid` | all | How width/height are produced: `sigmoid` or `exp`. Must match the checkpoint. |
 | `--epochs` | `500` | train | Number of epochs. |
 | `--batch-size` | `64` | train, sample, benchmark-workers | Batch size. |
-| `--learning-rate` | `1e-3` | train | Adam learning rate. |
-| `--weight-decay` | `1e-4` | train | Adam weight decay. |
+| `--learning-rate` | `1e-3` | train | Peak AdamW learning rate (reached after the warmup). |
+| `--weight-decay` | `0.05` | train | AdamW weight decay, applied to conv weights only (not biases or BatchNorm). |
+| `--warmup-epochs` | `3` | train | Linear warmup from `learning-rate / 10` to `learning-rate`. `0` disables it. |
+| `--min-learning-rate` | `1e-5` | train | Cosine decay ends at this value in the last epoch. |
 | `--train-workers` | `2` | train | DataLoader workers for training. |
 | `--val-workers` | `0` | train | DataLoader workers for validation. |
 | `--device` | `auto` | train, sample, benchmark-workers, webcam | `auto`, `cpu`, `cuda:0`, … |
@@ -341,8 +346,13 @@ Code: `src/yolo/` · Defaults: `src/yolo/config.py` · Model: [`Yolo_model`](#yo
 
 Trains the detector on COCO-format data.
 
-- **Optimizer and loss:** Adam with [`YoloLoss`](#loss).
-- **Best checkpoint:** chosen by **validation loss** (lower is better).
+- **Optimizer and loss:** AdamW with [`YoloLoss`](#loss). Weight decay only applies to
+  conv weights; biases and BatchNorm parameters are not decayed.
+- **Learning rate:** linear warmup for `--warmup-epochs`, then cosine decay from
+  `--learning-rate` to `--min-learning-rate` at the last epoch. The `lr` column
+  of the run log shows the value each epoch trained with.
+- **Best checkpoint:** chosen by **validation `conf_ap`** (higher is better).
+  Validation loss tends to keep falling after the detection quality peaks.
 - **Extra validation metrics:** after every epoch it also computes
   [detection metrics](#validation-metrics) on the whole validation set:
   `conf_ap`, `box_iou`, `precision` and `recall`.
@@ -418,6 +428,39 @@ python -m src.main yolo.webcam --camera 1 --threshold 0.9
 Needs an OpenCV build with GUI support (`opencv-python`, not the `-headless`
 variant from `requirements.txt`).
 
+#### `yolo.convert-wider`
+
+Converts the [WIDER FACE](http://shuoyang1213.me/WIDERFACE/) dataset into the
+COCO folders `yolo.train` 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, so it is not converted. The images are copied into each folder
+(flat) next to `_annotations.coco.json`, with one category, 1 = face.
+
+`--wider-dir` must contain the unpacked `WIDER_train.zip`, `WIDER_val.zip`
+and `wider_face_split.zip`. The files are found anywhere below it, so nested
+unzips like `WIDER_train/WIDER_train/images` also work.
+
+WIDER is mostly crowds of tiny faces, so the conversion filters for what the
+model can learn from:
+
+- **Small faces:** faces narrower or shorter than `--wider-min-face` px
+  after resizing to `--image-size` are dropped. Faces marked `invalid` are
+  dropped as well.
+- **Shared cells:** an image is skipped when two face centres fall in the
+  same `--grid` cell. The dataset keeps one face per cell, so the other face
+  would be trained as background.
+- **Boxes:** clipped to the image and stored in whole pixels.
+- **Unreadable images:** images OpenCV cannot read are skipped.
+
+Convert with the same `--image-size` and `--grid` you train with. The task
+refuses to overwrite existing `train/` or `test/` folders.
+
+```bash
+python -m src.main yolo.convert-wider --wider-dir ~/Downloads/wider_face
+python -m src.main yolo.convert-wider --wider-dir ~/Downloads/wider_face --image-size 128 --grid 8 --wider-out-dir data/wider_face_128
+python -m src.main yolo.train --train-dir data/wider_face/train --val-dir data/wider_face/test
+```
+
 ---
 
 ### `rcnn.*`: R-CNN-style experiments
@@ -843,11 +886,15 @@ far slower than the single-shot YOLO model. That is why the website uses YOLO.
    does not depend on the batch size. Optional extra metrics
    (`val_metrics`, e.g. YOLO's detection metrics) are computed the same way.
 3. **Compare:** checks the monitored metric against the best so far.
-   - `monitor="loss"`: lower is better (YOLO, R-CNN).
-   - any other metric, e.g. `"acc"`: higher is better (bird).
+   - `monitor="loss"`: lower is better (R-CNN).
+   - any other metric, e.g. `"acc"` (bird) or `"conf_ap"` (YOLO): higher is better.
 4. **Save on improvement:** if it improved, saves the checkpoint (`--save`)
    and the run's `best.pt`.
 5. **Log:** prints one summary line and appends a row to the run log.
+6. **Schedule:** steps the optional `scheduler` (YOLO uses `warmup_cosine`).
+
+`training.py` also has the optimizer helpers YOLO uses: `adamw(model, lr,
+weight_decay)` (no decay on biases and BatchNorm) and `warmup_cosine(...)`.
 
 Tasks plug in their own batch handling with `prepare_batch`. For example,
 `rcnn.train-football` cuts crops out of full images there.
@@ -1196,6 +1243,7 @@ folder. Unpack it in `api/` with `7z x data.7z`.
 | `faces_scenery/train/`, `test/` | COCO | `yolo.*` (default) | Face detection set: 10,825 training and 793 test images, category 1 = face. |
 | `faces/train/` | COCO | `rcnn.create-dataset` | Face set (3,910 images). |
 | `faces_processed/face/`, `background/` | ImageFolder | `rcnn.train-faces` | Crops made by `rcnn.create-dataset`. |
+| `wider_face/train/`, `test/` | COCO | `yolo.*` | Created by `yolo.convert-wider` from WIDER FACE, category 1 = face. Not in `data.7z`. |
 | `football/train/`, `valid/` | COCO | `rcnn.*-football` | Roboflow football players: 1 = ball, 2 = goalkeeper, 3 = player, 4 = referee. |
 | `faces_2/`, `mnist/`, WIDER Face archives | various | none | Not used by the current code. |
 

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,6 +26,30 @@ def _default_prepare(batch, device):
     return inputs.to(device), targets.to(device)
 
 
+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
@@ -84,6 +111,7 @@ def fit(
     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:
@@ -91,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, val_metrics, run)
+                           save_dir, model_name, save, monitor, track_accuracy, prepare_batch, val_metrics, run,
+                           scheduler)
     except KeyboardInterrupt:
         if run:
             run.finish("interrupted")
@@ -106,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, val_metrics, 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):
@@ -126,6 +156,8 @@ 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:

api/src/models/yolo.py

@@ -51,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),
@@ -60,6 +61,7 @@ class Yolo_model(nn.Module):
             SkipBlock(c_in=c_hidden, c_out=c_hidden),
             SkipBlock(c_in=c_hidden, c_out=c_hidden),
 
+            #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),
 

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/tasks.py

@@ -10,13 +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):
@@ -51,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)
@@ -59,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", val_metrics=detection_metrics, config=cfg)
+        monitor="conf_ap", val_metrics=detection_metrics, config=cfg, scheduler=scheduler)
 
 
 def view_data(cfg: Config):
@@ -147,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

@@ -2,9 +2,10 @@
 
 
 for ((c = 1; c <= 32; c *= 2)); do
-    python -m src.main yolo.train --train-dir ./data/faces_scenery/train/ --val-dir ./data/faces_scenery/test/ \
-        --no-augment-val --c-hidden "$c" --batch-size 64 --epochs 200 --run-name "bn-weighted-loss-c$c" \
-        --no-save --train-workers 4 --image-size 64 --grid 6
+    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