File size: 4,546 Bytes
e2f3b24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """
Train YOLO26 semantic segmentation on CelebAMask-HQ (19-class face parsing).
Variants: nano (realtime, ships in the iOS app) and large (accuracy reference).
imgsz=512 == deploy resolution (masks are natively 512). ENet inverse-log class
weighting is automatic in the semantic trainer (rare classes: glasses/hat/
earring/necklace).
Augmentation notes:
* fliplr=0.0 -- CRITICAL. RandomFlip flips geometry WITHOUT swapping the
left/right class labels, so any hflip corrupts l_eye/r_eye, l_brow/r_brow,
l_ear/r_ear. Must stay disabled.
* mosaic + degrees/translate/scale -- scale/position/tilt robustness to close
the aligned-crop (CelebA-HQ) -> live-selfie domain gap.
Throughput: uses MPSWorkersTrainer to undo ultralytics' forced workers=0 on MPS.
Timing probe: ml/train.py --epochs 1 --fraction 0.04 --no-val --name bench_X
Full run: ml/train.py --model yolo26l-sem.pt --epochs 60 --name celeba_large
"""
import argparse
from ultralytics import YOLO
from fast_trainer import MPSWorkersTrainer
DATA = "/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/data.yaml"
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--data", default=DATA)
ap.add_argument("--model", default="yolo26n-sem.pt")
ap.add_argument("--epochs", type=int, default=100)
ap.add_argument("--imgsz", type=int, default=512)
ap.add_argument("--batch", type=int, default=32)
ap.add_argument("--fraction", type=float, default=1.0)
ap.add_argument("--device", default="mps")
ap.add_argument("--workers", type=int, default=8)
ap.add_argument("--name", default="celeba_face")
ap.add_argument("--project", default="/Users/ari/FaceSegmentation/runs_semantic")
ap.add_argument("--patience", type=int, default=30)
ap.add_argument("--cache", default="False")
ap.add_argument("--mosaic", type=float, default=0.5)
ap.add_argument("--val", action=argparse.BooleanOptionalAction, default=True)
ap.add_argument("--plots", action=argparse.BooleanOptionalAction, default=True)
# deterministic=True (ultralytics default) forces slow MPS kernels in the
# semantic loss -- measured 5037 -> 2293 ms/iter when disabled.
ap.add_argument("--deterministic", action=argparse.BooleanOptionalAction, default=False)
ap.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
ap.add_argument("--val-period", type=int, default=1,
help="validate every Nth epoch (plus the final epoch)")
# ENet weights are (1/ln(1.02+p))**cls_pw. Ultralytics defaults cls_pw=0.0,
# which makes every weight exactly 1.0 -- i.e. class weighting silently OFF.
# At 1.0 the rare classes (necklace/eyes/glasses/earring) get ~45-50x the
# weight of hair/background/skin (~3.4-4.2), a 14.6x dynamic range.
ap.add_argument("--cls-pw", type=float, default=1.0,
help="ENet class-weight exponent; 0 disables class weighting")
# The stock loss upsamples stride-8 logits to full mask resolution (159M
# elements at 512/b32). Evaluating it lower is much cheaper but gives up
# some sub-grid boundary supervision. 0 = stock behaviour.
ap.add_argument("--loss-size", type=int, default=0,
help="compute the loss at this resolution instead of full mask res")
args = ap.parse_args()
if args.loss_size:
import loss_patch
loss_patch.apply(args.loss_size)
print(f"loss_patch: computing semantic loss at {args.loss_size}px")
MPSWorkersTrainer.val_period = args.val_period
model = YOLO(args.model)
model.train(
trainer=MPSWorkersTrainer,
data=args.data,
epochs=args.epochs,
imgsz=args.imgsz,
batch=args.batch,
device=args.device,
workers=args.workers,
project=args.project,
name=args.name,
fraction=args.fraction,
patience=args.patience,
cache=args.cache,
deterministic=args.deterministic,
amp=args.amp,
cls_pw=args.cls_pw,
# --- augmentation tuned for face parsing ---
fliplr=0.0, # CRITICAL: preserves left/right class semantics
flipud=0.0,
degrees=10.0, # head tilt
translate=0.1,
scale=0.5,
mosaic=args.mosaic,
close_mosaic=10,
mixup=0.0,
copy_paste=0.0,
# default hsv_* kept for lighting robustness
val=args.val,
plots=args.plots,
seed=0,
exist_ok=True,
)
if __name__ == "__main__":
main()
|