File size: 4,364 Bytes
9f7ad84 | 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 109 110 | """Train YOLOv8x with experimental settings on 100% data."""
import sys
import torch
import json
import shutil
from pathlib import Path
from ultralytics import YOLO
_orig = torch.load
def _safe(*a, **kw): kw["weights_only"] = False; return _orig(*a, **kw)
torch.load = _safe
# Parse args: seed imgsz epochs optimizer lr experiment_name [extra_args...]
SEED = int(sys.argv[1]) if len(sys.argv) > 1 else 42
IMGSZ = int(sys.argv[2]) if len(sys.argv) > 2 else 1280
EPOCHS = int(sys.argv[3]) if len(sys.argv) > 3 else 100
OPTIMIZER = sys.argv[4] if len(sys.argv) > 4 else "AdamW"
LR = float(sys.argv[5]) if len(sys.argv) > 5 else 0.0005
EXP_NAME = sys.argv[6] if len(sys.argv) > 6 else "default"
# Experiment presets
PRESETS = {
"light_aug": {"mosaic": 0.3, "copy_paste": 0.0, "mixup": 0.0, "degrees": 5, "scale": 0.5, "cls": 4.0},
"no_aug": {"mosaic": 0.0, "copy_paste": 0.0, "mixup": 0.0, "degrees": 0, "scale": 0.0, "cls": 4.0},
"high_cls": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 8.0},
"freeze10": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 4.0, "freeze": 10},
"cosine_low": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 4.0, "cos_lr": True, "lr0": 0.0001},
"default": {"mosaic": 1.0, "copy_paste": 0.3, "mixup": 0.2, "degrees": 10, "scale": 0.9, "cls": 4.0},
}
preset = PRESETS.get(EXP_NAME, PRESETS["default"])
print(f"Experiment: {EXP_NAME}")
print(f"Settings: seed={SEED} imgsz={IMGSZ} epochs={EPOCHS} opt={OPTIMIZER} lr={LR}")
print(f"Preset: {preset}")
work_dir = Path(f"train_exp_{EXP_NAME}")
work_dir.mkdir(exist_ok=True)
ann = json.load(open("input/train/annotations.json"))
categories = ann["categories"]
images = ann["images"]
annotations = ann["annotations"]
train_img_dir = work_dir / "images" / "train"
val_img_dir = work_dir / "images" / "val"
train_lbl_dir = work_dir / "labels" / "train"
val_lbl_dir = work_dir / "labels" / "val"
for d in [train_img_dir, val_img_dir, train_lbl_dir, val_lbl_dir]:
d.mkdir(parents=True, exist_ok=True)
img_anns = {}
for a in annotations:
img_anns.setdefault(a["image_id"], []).append(a)
for img in images:
src = Path("input/train/images") / img["file_name"]
if not src.exists(): continue
dst = train_img_dir / img["file_name"]
if not dst.exists(): shutil.copy2(src, dst)
iw, ih = img["width"], img["height"]
label_lines = []
for a in img_anns.get(img["id"], []):
x, y, w, h = a["bbox"]
cx = max(0, min(1, (x + w/2) / iw))
cy = max(0, min(1, (y + h/2) / ih))
nw = max(0, min(1, w / iw))
nh = max(0, min(1, h / ih))
label_lines.append(f"{a['category_id']} {cx} {cy} {nw} {nh}")
lbl_name = img["file_name"].rsplit(".", 1)[0] + ".txt"
(train_lbl_dir / lbl_name).write_text("\n".join(label_lines))
for img in images[:5]:
src = Path("input/train/images") / img["file_name"]
dst = val_img_dir / img["file_name"]
if src.exists() and not dst.exists():
shutil.copy2(src, dst)
lbl_name = img["file_name"].rsplit(".", 1)[0] + ".txt"
lbl_src = train_lbl_dir / lbl_name
if lbl_src.exists(): shutil.copy2(lbl_src, val_lbl_dir / lbl_name)
nc = len(categories)
cat_names = {c["id"]: c["name"] for c in categories}
names_list = [cat_names.get(i, f"class_{i}") for i in range(nc)]
data_yaml = work_dir / "data.yaml"
data_yaml.write_text(f"path: {work_dir.resolve()}\ntrain: images/train\nval: images/val\nnc: {nc}\nnames: {names_list}\n")
train_args = dict(
data=str(data_yaml), epochs=EPOCHS, imgsz=IMGSZ,
batch=2 if IMGSZ <= 1280 else 1, workers=0,
device=0 if torch.cuda.is_available() else "cpu",
seed=SEED, close_mosaic=10,
mosaic=preset.get("mosaic", 1.0),
copy_paste=preset.get("copy_paste", 0.3),
mixup=preset.get("mixup", 0.2),
degrees=preset.get("degrees", 10),
translate=0.2, scale=preset.get("scale", 0.9),
fliplr=0.0, optimizer=OPTIMIZER,
lr0=preset.get("lr0", LR), lrf=0.01,
warmup_epochs=5, cls=preset.get("cls", 4.0),
label_smoothing=0.1, save=True, save_period=25,
cos_lr=preset.get("cos_lr", False),
)
if "freeze" in preset:
train_args["freeze"] = preset["freeze"]
model = YOLO("yolov8x.pt")
model.train(**train_args)
print(f"Done! Experiment: {EXP_NAME}")
|