YOLO视觉检测应用入门篇（四）

YOLO视觉检测应用入门篇（四）
上一篇，我们对dataset有了一定的了解。
这一篇，是时候来训练一个模型了。

准备

我们假定你已经准备好了数据集和yolo环境，可以在yolo11n.pt的基础上进行进一步的训练。

训练

下面是示例代码

"""
YOLOv11增量训练完整脚本
功能：基于已有权重继续训练，集成增强数据与早停机制
"""

import os
import torch
from ultralytics import YOLO
from ultralytics.utils.torch_utils import de_parallel
from pathlib import Path

# 配置区 ================================================
class RetrainConfig:
    # 模型参数
    pretrained_weights = r"path\to\pt.pt"  # 初始权重路径
    # cbam_model= r"path\to\yolo11_CBAM.yaml" # 个性化的模型，这里先注释掉了    
    output_dir = r'path\to\output'     # 输出目录
    
    # 数据参数
    data_yaml = r"path\to\dataset.yaml"  # 数据集配置文件
    img_size = 640                 # 输入分辨率
    batch_size = 4                # 根据GPU调整(8/16/32)
    
    # 训练参数
    epochs = 300                   # 最大训练轮次
    initial_lr = 0.001             # 初始学习率
    optimizer_type = 'AdamW'       # 优化器选项(SGD, Adam, AdamW)
    weight_decay = 0.0005          # L2正则化强度
    
    # 增强参数
    augment = True                 # 启用增强配置
    mosaic_prob = 0.5              # 马赛克增强概率
    mixup_prob = 0.2               # MixUp概率
    amp = True,         # 启用混合精度
    
    # 早停参数
    patience = 5                  # 早停等待轮次
    min_delta = 0.001              # 视为改进的最小精度变化
    #freeze = [f"model.{x}.weight" for x in range(20)]  # 个性化模型用于冻结主干网络，这里注释掉了
    
    # 设备配置
    device = 'cuda' if torch.cuda.is_available() else 'cpu' # 有CUDA使用cuda，没有就cpu

# 训练执行 ==============================================
def main(cfg):
    # 确保输出目录存在
    os.makedirs(cfg.output_dir, exist_ok=True)
    
    # 加载预训练模型
    model = YOLO(cfg.pretrained_weights)
    #model = YOLO(cfg.cbam_model).load(cfg.pretrained_weights)
    #model = YOLO(cfg.cbam_model)
    # 上面两个适用于个性化模型，这里注释掉了
    # 修改模型配置增强参数
    model.augment = {
        'mosaic': cfg.mosaic_prob,
        'mixup': cfg.mixup_prob,
        'degrees': 45.0,    # 旋转角度范围
        'translate': 0.2,   # 平移幅度
        'shear': 0.3        # 错切幅度
    }
    
    # 构建训练参数
    train_args = {
        'data': cfg.data_yaml,
        'epochs': cfg.epochs,
        'batch': cfg.batch_size,
        'imgsz': cfg.img_size,
        'resume': False,          # 继续训练模式
        'patience': cfg.patience,
        'augment': cfg.augment,
        'lr0': cfg.initial_lr,
        'optimizer': cfg.optimizer_type,
        'weight_decay': cfg.weight_decay,
        'project': cfg.output_dir,
        'device': cfg.device,
        #'freeze':cfg.freeze,
        'amp':cfg.amp,
        'exist_ok': True
    }
    
    # 执行训练
    results = model.train(**train_args)
    

    print(f'✅ 训练完成！最佳权重已保存至: {cfg.output_dir}')

if __name__ == '__main__':
    config = RetrainConfig()
    main(config)

上述代码是一个使用YOLO进行训练的示例。
更加详细进行YOLO训练的参数，可以从YOLO Gitub官方获得。
也欢迎使用我Github的开源项目 YOLO Train GUI进行训练。

小结

这一篇，我们简单介绍了下用代码来进行Yolo模型训练的过程。
相信对于Yolo的应用理解又进了一步。
至此，我相信对于YOLO，已经有了较为入门的理解了。
虽然val和perdict并没有介绍，但我相信入门之后理解也会更为迅速。
下一篇，我们会对整个入门篇进行一个小结并简单介绍下后续的一些方向。