DetectionModelTraining/02_train.bat

91 lines
2.2 KiB
Batchfile

@echo off
chcp 65001 >nul
cls
:: 设置 Python 3.11 路径
set "PATH=C:\Users\Tellme\AppData\Local\Programs\Python\Python311\Scripts;C:\Users\Tellme\AppData\Local\Programs\Python\Python311;%PATH%"
echo ============================================================
echo 训练鞋子检测模型 (YOLOv8 + 640x640)
echo ============================================================
echo.
:: 设置数据集路径
set DATASET=datasets/openimages-shoes-yolo/data.yaml
:: 检查数据集是否存在
if not exist %DATASET% (
echo [错误] 找不到数据集配置文件: %DATASET%
echo.
echo 请先下载数据集:
echo python 01_download_dataset.py --source openimages --max-samples 5000
pause
exit /b 1
)
echo [信息] 数据集: %DATASET%
echo.
:: 选择模型
echo 选择模型:
echo 1. YOLOv8n (轻量级, 速度快)
echo 2. YOLOv8s (推荐, 速度和精度平衡)
echo 3. YOLOv8m (高精度, 较慢)
echo.
set /p MODEL_CHOICE="输入选择 (1-3, 默认 2): "
if "%MODEL_CHOICE%"=="" set MODEL_CHOICE=2
if "%MODEL_CHOICE%"=="1" (
set MODEL=yolov8n.pt
set DESC=YOLOv8n
)
if "%MODEL_CHOICE%"=="2" (
set MODEL=yolov8s.pt
set DESC=YOLOv8s (推荐)
)
if "%MODEL_CHOICE%"=="3" (
set MODEL=yolov8m.pt
set DESC=YOLOv8m
)
echo.
echo [信息] 使用模型: %DESC%
echo.
:: 训练参数
set EPOCHS=150
set IMGSZ=640
set BATCH=16
echo 训练参数:
echo - Epochs: %EPOCHS%
echo - Image Size: %IMGSZ%x%IMGSZ%
echo - Batch Size: %BATCH%
echo - Device: GPU (cuda:0)
echo.
echo ============================================================
echo 开始训练
echo ============================================================
echo.
yolo detect train data=%DATASET% model=%MODEL% epochs=%EPOCHS% imgsz=%IMGSZ% batch=%BATCH% device=0
if %ERRORLEVEL% neq 0 (
echo.
echo [错误] 训练失败!
pause
exit /b 1
)
echo.
echo ============================================================
echo 训练完成!
echo ============================================================
echo.
echo 模型保存在: runs/detect/train/weights/best.pt
echo.
echo 下一步: 运行 03_export_onnx.bat 导出 ONNX
echo.
pause