70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
"""
|
|
YOLOv8n 目标检测器训练脚本
|
|
基于准备好的 YOLO 数据集微调极轻量级模型，用于高精召回 672x480 背景图中的全部图标
|
|
"""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
import torch
|
|
from ultralytics import YOLO
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
logger = logging.getLogger("train_detector")
|
|
|
|
DATASET_YAML = Path("tcaptcha/data/yolo/dataset.yaml")
|
|
MODEL_DIR = Path("tcaptcha/models")
|
|
|
|
|
|
def train(epochs: int = 40, batch: int = 8, imgsz: int = 640):
|
|
if not DATASET_YAML.exists():
|
|
logger.error("数据集配置文件未找到: %s，请先运行 prepare_yolo.py", DATASET_YAML)
|
|
return
|
|
|
|
MODEL_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 选择加速设备
|
|
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
|
logger.info("采用训练加速设备: %s", device)
|
|
|
|
# 加载预训练 YOLOv8n (极轻量，只有 3.2M 参数)
|
|
model = YOLO("yolov8n.pt")
|
|
|
|
logger.info("开始微调训练，轮数: %d, 批大小: %d, 分辨率: %d", epochs, batch, imgsz)
|
|
|
|
results = model.train(
|
|
data=str(DATASET_YAML.resolve()),
|
|
epochs=epochs,
|
|
batch=batch,
|
|
imgsz=imgsz,
|
|
device=device,
|
|
project=str(MODEL_DIR / "runs"),
|
|
name="detector_exp",
|
|
exist_ok=True,
|
|
workers=2,
|
|
augment=True,
|
|
degrees=45.0, # 随机旋转增强抗噪能力
|
|
flipud=0.2,
|
|
fliplr=0.5,
|
|
mosaic=1.0,
|
|
verbose=True,
|
|
)
|
|
|
|
# 保存最佳权重到 models/yolo_detector.pt
|
|
best_pt = MODEL_DIR / "runs" / "detector_exp" / "weights" / "best.pt"
|
|
target_pt = MODEL_DIR / "yolo_detector.pt"
|
|
if best_pt.exists():
|
|
import shutil
|
|
shutil.copy(best_pt, target_pt)
|
|
logger.info("✓ 训练完成！最佳模型权重已保存至: %s", target_pt)
|
|
else:
|
|
logger.warning("未找到 best.pt 权重")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--epochs", type=int, default=35, help="训练轮数 (默认: 35)")
|
|
parser.add_argument("--batch", type=int, default=8, help="Batch size (默认: 8)")
|
|
args = parser.parse_args()
|
|
train(epochs=args.epochs, batch=args.batch)
|