train_from_checkpoint.py

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

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()