登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 业界 ›从零开始：基于CUDA 12.6的YOLOv5模型训练实战（RTX 205 ...

从零开始：基于CUDA 12.6的YOLOv5模型训练实战（RTX 2050显卡全流程）

别萧玉 2025-6-5 10:04:47
基于cuda12.6训练yolov5模型

前面完成了使用CPU调用yolov5s模型进行识别车辆，现在想训练自己的模型进行目标识别，使用CPU效率太低，尝试使用GPU加速的Pytorch，再重新整理了一下完整的流程

  • 注：

    • 显卡：RTX 2050
    • cuda：NVIDIA CUDA 12.7.33

  • 最后完成了一个识别doro头像的模型：
    1.png

2.gif

GitHub开源地址：https://github.com/ChengYull/YOLOv5-CUDA12_6-Training
环境部署

Anaconda环境

Anaconda 是一个用于数据科学和机器学习的开源 Python 发行版本，它简化了包管理和环境管理。可以创建隔离的 Python 环境，避免包冲突。

  • 隔离的Python环境非常重要，因为在训练时用到的图片标注工具需要Python3.9版本，而进行训练时用的是Python3.10版本
前往官网下载最新版即可：https://www.anaconda.com/download/
安装完成后使用win+r输入cmd打开命令行，输入conda --version命令
如果成功获取到了版本，说明安装成功
3.png

如果提示找不到conda命令，则需要手动配置下环境变量
找到系统变量中的Path，然后下Anaconda安装目录中的condabin目录添加进去
4.png

5.png

确定关闭环境变量窗口后，新开命令行重新输入conda --version命令 应当是能获取到对应安装的版本
至此Anaconda环境配置完成

  • 常用命令

    • 创建环境：conda create -n 环境名 python=版本号
    • 激活环境：conda activate 环境名
    • 删除环境：conda env remove --name 环境名

环境配置完成，顺便创建虚拟环境
命令行输入：
  1. conda create -n doro python=3.10
复制代码
创建环境
6.png

创建完成，在conda根目录的env文件夹下也会生成对应环境名的文件夹
7.png

激活环境
  1. conda activate doro
复制代码
8.png

Pytorch环境

这里已经决定使用GPU，需要查看显卡的cuda版本，安装对应支持的cuda版本，再安装对应的Pytorch版本
cuda安装

打开英伟达控制面板，电脑任务栏搜索可以直接搜索NVIDIA Contol Panel打开
9.png

找到系统信息
10.png

找到cuda版本信息
11.png

这里我的版本是12.7.33，也就是安装的cuda版本要小于等于12.7.33
打开cuda的下载地址：https://developer.nvidia.com/cuda-toolkit-archive
这里我选择了12.6.0版本
12.png

这里选择的网络版下载
13.png

14.png

注意可以选择自定义（精简默认安装C盘）
15.png

第一次安装尽量全选
16.png

等待下载安装
17.png

安装完成后，查看环境变量Path，确定安装完成
18.png

打开命令行，运行nvcc --version即可查看版本号
19.png

安装成功
cuDNN安装

cuDNN（CUDA Deep Neural Network Library）是 NVIDIA 官方推出的 GPU 加速深度学习库，专为深度神经网络（DNN）计算优化，支持 CNN（卷积神经网络）、RNN（循环神经网络）、Transformer 等模型的训练和推理加速。
cuDNN需要登录英伟达账户才能下载，没有直接注册即可
下载地址：https://developer.nvidia.com/rdp/cudnn-download
20.png

同上选择自定义 全选
21.png

记住安装目录
22.png

安装完成后打开对应目录
23.png

将bin、include、lib文件夹拷贝到cuda的安装目录下
24.png

查看环境变量中是否配置了这4项（若没有则补上）
25.png

cuDNN安装完成
Pytorch安装

打开Pytorch官网：https://pytorch.org/
往下滑找到下图模块，选择对应的版本（cuda12.6不支持Conda安装 仅能使用pip安装）复制安装命令
26.png

cuda12.6版本的Pytorch，清华镜像源暂时没更新，使用清华镜像会下载cpu版本，所以只能使用源地址下载，大概要下载两个小时
注意要在激活对应的Conda环境下运行命令
27.png

等待下载完成
完成后可以前往Pycharm中验证是否能够调用到GPU
新建项目，选择基础conda（在自定义环境中conda无法选取我们创建的安装Pytorch的虚拟环境）
28.png

创建后，进入设置-项目-Python解释器，修改环境为安装Pytorch的环境
29.png

创建程序 验证
代码：
  1. import torch
  3. flag = torch.cuda.is_available()
  4. print(flag)
  6. ngpu = 1
  7. # Decide which device we want to run on
  8. device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
  9. print(device)
  10. print(torch.cuda.get_device_name(0))
  11. print(torch.rand(3, 3).cuda())
复制代码
成功调用到GPU
30.png

至此所有环境配置完成
训练集准备

labelImg图片标注工具安装

给图片标注构造训练集，用于训练模型，由于在Python3.10环境下，labelImg软件存在闪退问题，这里使用Conda为LabelImg创建一个新的Python3.9环境
  1. conda create -n label python=3.9
复制代码
注意先激活环境再安装LabelImg
  1. conda activate label
复制代码
安装LabelImg
  1. conda install pyqt labelimg
复制代码
31.png

找到创建的环境目录，找到Scrip文件夹
32.png

33.png

打开界面如下
34.png

可以参考一下我的设置
35.png

注意：save下方一定要选择yolo模型，否则标注保存的文件是xml文件而非txt文件
标注图片

创建一个单独的文件夹，我这里命名train，再在里面创建两个文件夹（images、labels）以及一个yaml文件（doro.yaml）
36.png

yaml文件内容：
  1. # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
  2. train: ../train/images/
  3. val: ../train/images/
  4. # number of classes
  5. nc: 1
  7. # class names
  8. names: ['doro']
复制代码
准备好用于训练的图片，我这里是在网上下载的表情包，大概有五十多张，放到刚刚创建的images文件夹中
37.png

打开LabelImg，选择打开文件夹
38.png

选择保存文件夹位置
39.png

设置一下label名称，就无需手动输入了
40.png

按键w可以快捷创建选框，按键A和D是切换上下张图片
完成后在labels文件夹下就会出现标注的内容
41.png

至此训练集准备完成
模型训练

yolov5模型源码部署

Github仓库获取yolov5源码：https://github.com/ultralytics/yolov5
这里我使用的是Pycharm克隆项目
42.png

43.png

将对应地址填入克隆配置中
44.png

克隆完成如图
45.png

同上设置Python解释器
46.png

安装依赖
  1. pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
复制代码
47.png

安装完成后测试运行（第一次运行会自动下载yolov5s.pt模型）
  1. python detect.py --weights yolov5s.pt --source data/images/bus.jpg
复制代码
48.png

在runs目录如果可以查看到运行结果，说明运行正常
49.png

训练

这里为了方便，将之前创建的train文件夹（存放训练集）移动到项目中
50.png

在目录中创建src文件夹，用于存放自己的代码
51.png

配置训练参数（数据集路径需要改为对应的yaml文件路径）
  1. from train import parse_opt, main
  2. import torch
  3. # 在训练前清理GPU缓存
  4. torch.cuda.empty_cache()
  5. def custom_train():
  6.     opt = parse_opt()
  7.     opt.data = "../train/doro.yaml"  # 数据集路径
  8.     opt.weights = "../yolov5s.pt"    # 权重文件
  9.     opt.imgsz = 640                # 图像尺寸
  10.     opt.batch_size = 8             # 批量大小
  11.     opt.epochs = 100               # 训练轮次
  12.     opt.nosave = False             # 允许保存检查点
  13.     opt.save_period = 20           # 每20轮保存一次
  14.     opt.cache = True               # 启用数据缓存
  15.     opt.device = "0"               # 使用GPU 0（若为CPU则设为"cpu"）
  16.     opt.lr0 = 0.01                 # 初始学习率
  18.     # 启动训练
  19.     main(opt)
  21. if __name__ == "__main__":
  22.     custom_train()
复制代码
其中图像尺寸imgsz和批量大小batch_size要根据显存配置，如果训练过程中出现torch.OutOfMemoryError: CUDA out of memory，试着将其适当调小
52.png

如果多次修改无果，可以试着重启电脑（亲测有效）
训练正式开始，等待100轮的训练
53.png

训练完成后，在runs文件夹下可以找到训练好的模型文件（如果不想保存中间生成的模型文件，可以在训练参数将opt.nosave = True改为True，删掉opt.save_period = 20配置）
54.png

测试模型

这里使用openCV来读取视频来识别作为测试
  1. import cv2
  2. import torch
  5. # 读取视频
  6. video_path = "../testVideo/doro3.mp4"
  7. cap = cv2.VideoCapture(video_path)
  9. # 加载训练的模型
  10. model = torch.hub.load('../../test_doro', 'custom', path='../runs/train/exp/weights/best.pt', source='local')
  12. # 检查视频是否成功打开
  13. if not cap.isOpened():
  14.     print("无法打开视频文件")
  15.     exit()
  17. # 播放视频
  18. while True:
  19.     ret, frame = cap.read()
  20.     if not ret:
  21.         break
  23.     # 模型推理
  24.     results = model(frame)
  25.     # 获取预测结果
  26.     for pred in results.pred[0]:
  27.         x1, y1, x2, y2, conf, cls = pred.tolist()
  28.         class_name = model.names[int(cls)]
  29.         # 输出结果
  30.         print(f"检测到：{class_name}, 置信度：{conf:.2f}")
  31.         cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
  32.         cv2.putText(frame, f"{class_name} {conf:.2f}", (int(x1), int(y1) - 10),
  33.                     cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
  34.     # 显示当前帧
  35.     cv2.imshow("Video", frame)
  36.     # 按下 'q' 键退出
  37.     if cv2.waitKey(1) & 0xFF == ord('q'):
  38.         break
  40. # 释放视频捕获对象和关闭所有窗口
  41. cap.release()
  42. cv2.destroyAllWindows()
复制代码
这里发现能够成功识别，且效果还行，但是发现控制台有一行警告信息，非常影响观感
55.png

根据提示，找到models目录下的common.py文件，搜索 with amp.autocast(autocast):
56.png

将其替换为：
  1. with torch.amp.autocast("cuda",enabled=autocast):
复制代码
这里有两处：分别在879行和906行
57.png

58.png

再次运行测试代码，警告消失
59.png


来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

使用道具 举报
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
别萧玉
2025-6-5 10:04:47

0

粉丝关注

11

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫