From 4b227e51fcda8400fb494d5cfa0670f050106747 Mon Sep 17 00:00:00 2001 From: Hamza Zakaria Date: Wed, 29 Apr 2026 09:00:32 +0100 Subject: [PATCH] train_from_checkpoint.py --- runs/detect/runs/detect/coffee_v1/args.yaml | 8 +- train_from_checkpoint.py | 181 ++++++++++++++++++++ 2 files changed, 185 insertions(+), 4 deletions(-) create mode 100644 train_from_checkpoint.py diff --git a/runs/detect/runs/detect/coffee_v1/args.yaml b/runs/detect/runs/detect/coffee_v1/args.yaml index acd1385..b84bd17 100644 --- a/runs/detect/runs/detect/coffee_v1/args.yaml +++ b/runs/detect/runs/detect/coffee_v1/args.yaml @@ -1,6 +1,6 @@ task: detect mode: train -model: yolov8n.pt +model: runs\detect\runs\detect\coffee_v1\weights\last.pt data: dataset.yml epochs: 150 time: null @@ -11,7 +11,7 @@ save: true save_period: 10 cache: false device: cpu -workers: 8 +workers: 0 project: runs/detect name: coffee_v1 exist_ok: true @@ -24,7 +24,7 @@ single_cls: false rect: false cos_lr: false close_mosaic: 10 -resume: false +resume: runs\detect\runs\detect\coffee_v1\weights\last.pt amp: true fraction: 1.0 profile: false @@ -77,7 +77,7 @@ momentum: 0.938 weight_decay: 0.0005 warmup_epochs: 5 warmup_momentum: 0.8 -warmup_bias_lr: 0.1 +warmup_bias_lr: 0.0 box: 7.5 cls: 0.5 cls_pw: 0.0 diff --git a/train_from_checkpoint.py b/train_from_checkpoint.py new file mode 100644 index 0000000..e4bb588 --- /dev/null +++ b/train_from_checkpoint.py @@ -0,0 +1,181 @@ +""" +train_from_checkpoint.py +======================== +Resume training from last saved checkpoint runs/detect/coffee_v1/weights/last.pt + +Usage: + python train_from_checkpoint.py + +Outputs: + runs/detect/coffee_v1/weights/best.pt <- use this for detection. +""" + +import os +import torch + +from pathlib import Path +from ultralytics import YOLO + +# ── Configuration ─────────────────────────────────────────────────── +os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" + +DATASET_YAML = "dataset.yml" # Generated by prepare_dataset.py + +PROJECT = "runs/detect" +RUN_NAME = "coffee_v1" +IMG_SIZE = 640 + +# ── Training hyperparameters ─────────────────────────────────────────────────── +# +# Tuned for: small dataset (few-shot), packaged goods, retail shelf. +# +EPOCHS = 150 +PATIENCE = 30 # Early Stopping - stops if no improvement +BATCH = 8 # Lower if you get OOM on CPU/small GPU. +LR0 = 0.005 # Initial LR (Learing Rate) - lower than default (0.01) + +LRF = 0.01 # final LR = LR0 * LRF +MOMENTUM = 0.938 +WEIGHT_DECAY = 0.0005 +WARMUP_EPOCHS = 5 # Longer warmup stabilises few-shot training. +IOU_THRESH = 0.5 # IoU threshold for NMS. +CONF_THRESH = 0.25 # confidence threshold at val time. + +# ── YOLOv8 built-in augmentation (on top of our albumentatios pass) ──────────── +# +# These run during tarining batches - complementary to augment.py. +# Mosaic is the single most impactful augmentation for small datasets. +# +MOSAIC = 1.0 # mosaic probability +MIXUP = 0.1 # mixup probability +COPY_PASTE = 0.1 +DEGREES = 5.0 # degrees ±5° +TRANSLATE = 0.1 +SCALE = 0.5 +FLIPUD = 0.0 +FLIPLR = 0.5 +HSV_H = 0.015 +HSV_S = 0.7 +HSV_V = 0.4 + + +def check_enviroment(): + device = "cuda" if torch.cuda.is_available() else "cpu" + if device == "cpu": + print("[INFO] No GPU detected - trainning on CPU") + print(" Expected time: ~60-90min for 150 epochs on a typical laptop.") + print(" Tips: use Google Colab or Kaggle they got free GPU, if training is too slow.\n") + else: + gpu= torch.cuda.get_device_name(0) + print(f"[INFO] GPU detected: {gpu}") + print(" Expected time: ~5-15min.\n") + return device + +def fix_grad_scaler(checkpoint_path: str): + """ + Patches a checkpoint whose GradScaler state is empty + (happens when last.pt was saved on CPU / AMP-disabled env). + """ + ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + + default_scaler_state = { + "scale": torch.tensor(65536.0), # standard initial scale + "growth_factor": 2.0, + "backoff_factor": 0.5, + "growth_interval": 2000, + "_growth_tracker": 0, + } + + patched = False + # Ultralytics stores the scaler under the key 'scaler' + if "scaler" in ckpt: + if not ckpt["scaler"]: # empty dict → broken + ckpt["scaler"] = default_scaler_state + patched = True + else: + ckpt["scaler"] = default_scaler_state + patched = True + + if patched: + torch.save(ckpt, checkpoint_path) + print(f"[INFO] GradScaler state patched in {checkpoint_path}") + else: + print(f"[INFO] GradScaler state looks fine, no patch needed.") + +def train(): + device = check_enviroment() + + if not Path(DATASET_YAML).exists(): + print(f"[ERROR] {DATASET_YAML} not found.") + print(" Run prepare_dataset.py first.") + return + + model_path = Path(PROJECT) / Path(PROJECT) / RUN_NAME / "weights" / "last.pt" + if not model_path.exists(): + print(f"[ERROR] {model_path} not found.") + print(" Run train.py first.") + return + + # Patch checkpoint before loading. + if device == "cuda": + fix_grad_scaler(str(model_path)) + + model = YOLO(model_path) + print(f"Starting training: {EPOCHS} epochs, img size {IMG_SIZE}") + print(f"Dataset: {DATASET_YAML}") + print(f"Output: {PROJECT}/{RUN_NAME}/weights/best.pt\n") + + results = model.train( + resume = True, + data = DATASET_YAML, + epochs = EPOCHS, + patience = PATIENCE, + imgsz = IMG_SIZE, + batch = BATCH, + device = device, + project = PROJECT, + name = RUN_NAME, + exist_ok = True, + + lr0 = LR0, + lrf = LRF, + momentum = MOMENTUM, + weight_decay = WEIGHT_DECAY, + warmup_epochs = WARMUP_EPOCHS, + + iou = IOU_THRESH, + conf = CONF_THRESH, + + # Built-in augmentation + mosaic = MOSAIC, + mixup = MIXUP, + copy_paste = COPY_PASTE, + degrees = DEGREES, + translate = TRANSLATE, + scale = SCALE, + flipud = FLIPUD, + fliplr = FLIPLR, + hsv_h = HSV_H, + hsv_s = HSV_S, + hsv_v = HSV_V, + + save = True, + save_period = 10, # Checkpoint every 10 epochs + plots = True, # Saves training curvse as PNGs + verbose = True, + seed = 42, + ) + + best_weights = Path(PROJECT) / RUN_NAME / "weights" / "best.pt" + print("\n" + "-" * 60) + print(f"Training complete.") + print(f"Best weights saved -> {best_weights}") + print(f"mAP50: {results.results_dict.get('metrics/mAP50(B)', 'N/A'):.3f}") + print(f"mAP50-95: {results.results_dict.get('metrics/mAP50-95(B)', 'N/A'):.3f}") + print("-" * 60) + print("\nRun next: python detect.py --source scene.jpg") + + return best_weights + +if __name__ == "__main__": + train()