Windows + PyTorch + YOLOv8：小白也能玩转AI目标检测

YOLO模型YOLO（You Only Look Once）是一种高效的目标检测模型，它将图像划分为网格，并在每个网格中同时预测边界框和类别，从而实现端到端的实时目标检测。与传统方法相比，YOLO具有速度快、精度高、适用于实时应用的优点，广泛应用于视频监控、自动驾驶等领域。

博主最近也在学习神经网络相关知识，本期记录一下Windows系统训练和部署YOLO神经网络模型的过程。

系统环境：Python 3.9、PyCharm 2022

1、环境安装

Python 3.9和PyCharm的环境安装不做过多的描写，我们主要介绍搭建一个Pytorch和YOLO环境搭建。

（1-1）首先搭建一个Pycharm的基本工程，Python释器版本选择3.9版本。

图2 PyCharm默认工程图

（1-2）打开终端我们安装需要对应PyTorch和YOLO的库文件。

图3   终端使用图

（ GPU版本）

PyTorch的模型训练可以选择CPU训练或者GPU加速，GPU加速的情况下训练速度会比CPU快很多很多很多！！！因此我们需要查询和自己GPU适配的PyTorch版本（本文内容适用于N卡）。

（1-3可选）查询CUDA版本

 nvidia-smi

查询CUDA 版本

图4   CUDA版本查看

可以看到我的CUDA版本为12.6。

（1-4）接下来根据CUDA版本我们使用pip命令来安装PyTorch,大家可以根据官网推荐的版本来安装。

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

安装cuda12.1版本的pytorch （12.6兼容12.1）

图5    Pytorch安装

库文件比较大，大概有2G+，如果遇到网络问题可以多安装几次。

（1-5）查询pytorch是否安装成功

python -c "import torch; print(torch.cuda.is_available())"

查询pytorch并且是否支持GPU

图6    pytorch验证

返回True即可代表我们可以正常使用GPU

（1-6）安装YOLO

pip install ultralytics

YOLO库代码安装

图7    Yolo验证

    安装成功后输入yolo会返回很多内容，接下来实现模型训练。

2、训练集准备

图8    文件存放架构图

（2-1）按照这个路径我们存放训练集和验证集，images和labels分别存储图像数据和标签。train文件夹存放用于训练的数据，val存放用于验证的数据。

labels是YOLO格式的txt文件如图所示：

图9    Yolo标注格式

    它表达了我们要识别的内容在图像中的位置信息，可以使用标注工具来制作（可以查询Labelimg的教学）

图10 Labelimg使用实拍

    （2-2）制作配置文件dataset.yaml（文件名不必相同）

train: ./data/images/train #训练集地址val: ./data/images/val    #验证集地址nc: 4 # 类别数量names: ['Chongdianbao', 'Red Tea', 'Pen', 'Upan']  # 类别名称

dataset.yaml配置文件用于等会模型训练

图11    yaml配置文件

    （2-3）模型训练代码

import torchfrom ultralytics import YOLOimport multiprocessing
def run():    model = YOLO('yolov8n.pt')  # 直接使用预训练模型微调    model.train(        data = './dataset.yaml',#模型配置文件        epochs = 800,#总训练轮数800轮        imgsz =  480,#图片缩放为480*480        batch = 20,  #一批20张        name = 'yolov8_custom',#模型名字        device = 0,  #使用GPU进行计算        val = True   #训练过程中执行验证      )    # 评估模型    results = model.val()    print(results)

if __name__ == '__main__':    multiprocessing.freeze_support()    run()

模型训练Demo

图12    训练时命令行截图

    接下来命令行会实时显示训练过程，等待训练结束后，可以在runs文件夹中找到我们的模型。

图13    结果图

图14    模型结果

    weights中存放着模型，其中best.pt为最优模型，last.pt为上一次模型。

（2-4）模型使用

from ultralytics import YOLOimport cv2# 1. 加载训练好的模型（可以是 yolov8n.pt 或你的 best.pt）model = YOLO("best.pt") # 2. 读取一张图片img = cv2.imread("test.jpg")# 3. 进行推理（返回结果列表）results = model.predict(source=img, save=False)# 4. 将检测框画在图上并显示annotated = results[0].plot()cv2.imshow("YOLO Detection", annotated)cv2.waitKey(0)cv2.destroyAllWindows()

模型使用Demo

3、测试效果

标注过程是真的痛苦啊，感觉像是电子打螺丝