下面我将为你提供一个完整的 YOLOv8n(nano 版本)从环境配置到训练再到推理的详细步骤和代码。YOLOv8 是 Ultralytics 推出的最新版本目标检测算法,nano 版本是其中最轻量化的模型,适合在资源有限的设备上运行。

一、环境配置

首先需要配置 Python 环境并安装必要的依赖库:

# 创建并激活虚拟环境
conda create -n yolov8 python=3.9 -y
conda activate yolov8

# 安装PyTorch (根据自己的CUDA版本调整,这里以CUDA 11.7为例)
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

# 安装Ultralytics库 (包含YOLOv8)
pip install ultralytics

# 安装其他依赖
pip install opencv-python matplotlib pandas seaborn tqdm

如果你没有 NVIDIA 显卡,可以使用 CPU 版本的 PyTorch:

pip3 install torch torchvision torchaudio

二、数据集准备

YOLOv8 需要特定格式的数据集,下面是准备数据集的步骤:

  1. 数据集结构如下:
dataset/
├── images/
│   ├── train/
│   │   ├── img1.jpg
│   │   ├── img2.jpg
│   │   └── ...
│   └── val/
│       ├── img1.jpg
│       ├── img2.jpg
│       └── ...
└── labels/
    ├── train/
    │   ├── img1.txt
    │   ├── img2.txt
    │   └── ...
    └── val/
        ├── img1.txt
        ├── img2.txt
        └── ...

2.创建数据集配置文件(dataset.yaml):

# 训练和验证集的路径
train: ./dataset/images/train
val: ./dataset/images/val

# 类别数量
nc: 80  # 以COCO数据集为例,有80个类别

# 类别名称
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
         'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
         'hair drier', 'toothbrush' ]

如果你没有自己的数据集,可以使用 COCO 数据集的子集进行训练,或者使用 Ultralytics 提供的示例数据集。

三、模型训练

下面是使用 YOLOv8n 进行模型训练的代码:

from ultralytics import YOLO
import argparse

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='YOLOv8n 模型训练')
    parser.add_argument('--data', type=str, default='dataset.yaml', help='数据集配置文件路径')
    parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
    parser.add_argument('--batch', type=int, default=16, help='批次大小')
    parser.add_argument('--imgsz', type=int, default=640, help='输入图像大小')
    parser.add_argument('--device', type=str, default='0', help='训练设备,0表示第一块GPU,cpu表示使用CPU')
    parser.add_argument('--name', type=str, default='yolov8n_custom', help='训练任务名称')
    parser.add_argument('--pretrained', action='store_true', help='是否使用预训练模型')
    return parser.parse_args()

def train_yolov8n(args):
    """训练YOLOv8n模型"""
    # 加载YOLOv8n模型
    if args.pretrained:
        # 使用预训练模型
        model = YOLO('yolov8n.pt')
    else:
        # 从零开始训练
        model = YOLO('yolov8n.yaml')
    
    # 开始训练
    results = model.train(
        data=args.data,
        epochs=args.epochs,
        batch=args.batch,
        imgsz=args.imgsz,
        device=args.device,
        name=args.name,
        # 其他可配置参数
        patience=10,  # 早停耐心值
        save=True,    # 保存模型
        save_period=10,  # 每10轮保存一次
        pretrained=args.pretrained,  # 是否使用预训练权重
        optimizer='Adam',  # 优化器
        lr0=0.001,    # 初始学习率
        augment=True, # 启用数据增强
        cache=True    # 缓存图像以加速训练
    )
    
    # 评估模型
    metrics = model.val()
    print("模型评估结果:")
    print(metrics)
    
    return model

if __name__ == '__main__':
    args = parse_args()
    model = train_yolov8n(args)
    print("训练完成!")

训练命令示例:

python train_yolov8n.py --data dataset.yaml --epochs 50 --batch 8 --imgsz 640 --device 0 --pretrained --name my_yolov8n

四、模型推理

训练完成后,可以使用以下代码进行模型推理:

    from ultralytics import YOLO
    import cv2
    import argparse
    import os
    from datetime import datetime
    
    def parse_args():
        """解析命令行参数"""
        parser = argparse.ArgumentParser(description='YOLOv8n 模型推理')
        parser.add_argument('--model', type=str, required=True, help='模型权重文件路径')
        parser.add_argument('--source', type=str, required=True, help='输入图像或视频路径,或0表示摄像头')
        parser.add_argument('--conf', type=float, default=0.25, help='置信度阈值')
        parser.add_argument('--iou', type=float, default=0.45, help='IOU阈值')
        parser.add_argument('--device', type=str, default='0', help='推理设备,0表示第一块GPU,cpu表示使用CPU')
        parser.add_argument('--save', action='store_true', help='是否保存结果')
        return parser.parse_args()
    
    def infer_image(model, image_path, conf, iou, save=False):
        """对单张图像进行推理"""
        # 加载图像
        img = cv2.imread(image_path)
        if img is None:
            print(f"无法加载图像: {image_path}")
            return
        
        # 进行推理
        results = model(img, conf=conf, iou=iou)
        
        # 处理结果
        for result in results:
            # 绘制检测框
            annotated_img = result.plot()
            
            # 显示结果
            cv2.imshow('YOLOv8 Detection', annotated_img)
            cv2.waitKey(0)
            
            # 保存结果
            if save:
                # 创建保存目录
                os.makedirs('inference_results', exist_ok=True)
                # 生成保存文件名
                filename = os.path.basename(image_path)
                name, ext = os.path.splitext(filename)
                save_path = f"inference_results/{name}_result{ext}"
                cv2.imwrite(save_path, annotated_img)
                print(f"结果已保存至: {save_path}")
        
        cv2.destroyAllWindows()
    
    def infer_video(model, video_path, conf, iou, save=False):
        """对视频进行推理"""
        # 打开视频
        cap = cv2.VideoCapture(video_path)
        
        # 获取视频信息
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # 视频写入器
        out = None
        if save:
            os.makedirs('inference_results', exist_ok=True)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_path = f"inference_results/video_result_{timestamp}.mp4"
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(save_path, fourcc, fps, (width, height))
            print(f"视频结果将保存至: {save_path}")
        
        # 逐帧处理
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            
            # 推理
            results = model(frame, conf=conf, iou=iou)
            
            # 绘制检测框
            annotated_frame = results[0].plot()
            
            # 显示结果
            cv2.imshow('YOLOv8 Detection', annotated_frame)
            
            # 保存帧
            if save and out is not None:
                out.write(annotated_frame)
            
            # 按q退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        # 释放资源
        cap.release()
        if save and out is not None:
            out.release()
        cv2.destroyAllWindows()
    
    def main(args):
        """主函数"""
        # 加载模型
        model = YOLO(args.model)
        model.to(args.device)
        
        # 判断输入类型并进行推理
        if args.source.isdigit():
            # 摄像头
            infer_video(model, int(args.source), args.conf, args.iou, args.save)
        elif args.source.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            # 图像文件
            infer_image(model, args.source, args.conf, args.iou, args.save)
        else:
            # 视频文件
            infer_video(model, args.source, args.conf, args.iou, args.save)
    
    if __name__ == '__main__':
        args = parse_args()
        main(args)
    

    四、模型推理

    训练完成后,可以使用以下代码进行模型推理:

    YOLOv8n模型推理代码

    推理命令示例:

    对单张图片进行推理:

    python infer_yolov8n.py --model runs/detect/my_yolov8n/weights/best.pt --source test_image.jpg --conf 0.3 --save
    

    对视频进行推理:

      python infer_yolov8n.py --model runs/detect/my_yolov8n/weights/best.pt --source test_video.mp4 --conf 0.3 --save
      

      使用摄像头进行实时推理:

      python infer_yolov8n.py --model runs/detect/my_yolov8n/weights/best.pt --source 0 --conf 0.3
      

      五、结果分析

      训练完成后,可以使用以下代码分析训练结果:

      import os
      import matplotlib.pyplot as plt
      import pandas as pd
      import seaborn as sns
      from ultralytics import YOLO
      
      def analyze_training_results(results_dir):
          """分析训练结果"""
          # 设置中文字体
          plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
          
          # 加载训练日志
          results_path = os.path.join(results_dir, "results.csv")
          if not os.path.exists(results_path):
              print(f"未找到训练日志: {results_path}")
              return
          
          results = pd.read_csv(results_path)
          
          # 创建图表保存目录
          os.makedirs(os.path.join(results_dir, "analysis"), exist_ok=True)
          
          # 1. 绘制损失函数曲线
          plt.figure(figsize=(12, 8))
          plt.plot(results['epoch'], results['train/box_loss'], label='训练集边界框损失')
          plt.plot(results['epoch'], results['train/cls_loss'], label='训练集分类损失')
          plt.plot(results['epoch'], results['train/dfl_loss'], label='训练集分布焦点损失')
          plt.plot(results['epoch'], results['val/box_loss'], label='验证集边界框损失', linestyle='--')
          plt.plot(results['epoch'], results['val/cls_loss'], label='验证集分类损失', linestyle='--')
          plt.plot(results['epoch'], results['val/dfl_loss'], label='验证集分布焦点损失', linestyle='--')
          plt.xlabel('轮次')
          plt.ylabel('损失值')
          plt.title('训练与验证损失曲线')
          plt.legend()
          plt.grid(True)
          plt.savefig(os.path.join(results_dir, "analysis", "loss_curves.png"))
          plt.close()
          
          # 2. 绘制mAP曲线
          plt.figure(figsize=(12, 8))
          plt.plot(results['epoch'], results['metrics/mAP50-95(B)'], label='mAP@50-95')
          plt.plot(results['epoch'], results['metrics/mAP50(B)'], label='mAP@50')
          plt.xlabel('轮次')
          plt.ylabel('mAP值')
          plt.title('mAP指标曲线')
          plt.legend()
          plt.grid(True)
          plt.savefig(os.path.join(results_dir, "analysis", "map_curves.png"))
          plt.close()
          
          # 3. 绘制精确率和召回率曲线
          plt.figure(figsize=(12, 8))
          plt.plot(results['epoch'], results['metrics/precision(B)'], label='精确率')
          plt.plot(results['epoch'], results['metrics/recall(B)'], label='召回率')
          plt.xlabel('轮次')
          plt.ylabel('值')
          plt.title('精确率和召回率曲线')
          plt.legend()
          plt.grid(True)
          plt.savefig(os.path.join(results_dir, "analysis", "precision_recall_curves.png"))
          plt.close()
          
          print(f"分析结果已保存至: {os.path.join(results_dir, 'analysis')}")
      
      def main():
          # 训练结果目录,根据实际情况修改
          results_dir = "runs/detect/my_yolov8n"
          
          # 分析训练结果
          analyze_training_results(results_dir)
          
          # 加载最佳模型并显示性能指标
          best_model_path = os.path.join(results_dir, "weights", "best.pt")
          if os.path.exists(best_model_path):
              model = YOLO(best_model_path)
              metrics = model.val()
              print("\n最佳模型性能指标:")
              print(f"mAP@50-95: {metrics.box.map:.4f}")
              print(f"mAP@50: {metrics.box.map50:.4f}")
              print(f"精确率: {metrics.box.mp:.4f}")
              print(f"召回率: {metrics.box.mr:.4f}")
          else:
              print(f"未找到最佳模型: {best_model_path}")
      
      if __name__ == "__main__":
          main()
      

      总结

      以上就是 YOLOv8n 从环境配置到训练再到推理的完整流程。通过这些代码,你可以:

      1. 配置 YOLOv8 所需的运行环境
      2. 准备符合 YOLO 格式的数据集
      3. 使用预训练权重或从零开始训练 YOLOv8n 模型
      4. 对图像、视频或摄像头进行实时目标检测
      5. 分析训练结果,评估模型性能

      根据你的具体任务和数据集,你可能需要调整一些参数,如训练轮数、批次大小、图像尺寸、置信度阈值等,以获得最佳的检测效果。

      Logo

      火山引擎开发者社区是火山引擎打造的AI技术生态平台,聚焦Agent与大模型开发,提供豆包系列模型(图像/视频/视觉)、智能分析与会话工具,并配套评测集、动手实验室及行业案例库。社区通过技术沙龙、挑战赛等活动促进开发者成长,新用户可领50万Tokens权益,助力构建智能应用。

      更多推荐