train_from_checkpoint.py
This commit is contained in:
parent
1afc0ea96c
commit
4b227e51fc
@ -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
|
||||
|
||||
181
train_from_checkpoint.py
Normal file
181
train_from_checkpoint.py
Normal file
@ -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()
|
||||
Loading…
Reference in New Issue
Block a user