【深度学习】使用GPU（CUDA）跑通YOLOv5源码，包含搭建环境+训练数据集+预测（整套流程）_使用gpu(cuda)

tips：本人是因为在CPU上训练数据太慢才选择使用GPU的，其实CPU和GPU和代码上只有一丝丝的区别。如果在搭建环境和跑通YOLOv5源码有遇到一些BUG可以参考我的上一篇文章： 【深度学习】浅浅记录卸载旧版本anaconda+安装pytorch环境+跑通YOLOv5源代码所遇到的一系列问题

---

1.安装软件

因为我是在全新的电脑上进行安装的，故需要自己进行pycharm和anaconda的安装。

1.1 anaconda安装

anaconda的安装可以参考我上一篇文章【深度学习】浅浅记录卸载旧版本anaconda+安装pytorch环境+跑通YOLOv5源代码所遇到的一系列问题里面的第二点，就是anaconda的安装教程。
简述步骤就是：①在清华源下载Anaconda3-2020.11-Windows-x86_64版本。②安装anaconda。③配置环境变量。④检查环境是否配置完成。

1.2 pycharm安装

此处我安装的是社区版，直接进入官网链接进行安装即可。
此处贴一个安装教程：程序员小麦的文章2023最新PyCharm安装+界面配置，给你一个舒适的开发环境
跟着此教程进行安装和汉化即可。

2.搭建环境

2.1 搭建anaconda环境

使用命令进行pytorch的虚拟环境安装，并进入创建好的ptyorch_GPU虚拟环境

conda create -n pytorch_GPU python=3.8

conda activate pytorch
1
2
3

2.2 查询CUDA版本号

CMD进入NVIDIA Corporation，再使用命令进行查看。
可以看到cudaversion可以兼容12.2及以下，也可以看到本机上显卡的编号是0（后面要用）

nvidia-smi
1

2.3 导入GPU版本的pytorch

去pytorch官网：https://pytorch.org/，选择相应的GPU版本。
因为兼容12.2及以下，就选择CUDA12.1即可。

复制最后一行的命令行，到上面创建好的虚拟环境pytorch_GPU中执行即可。
然后就是等待安装即可。
安装完成：

检查是否安装成功：

3.跑通YOLOv5项目

3.1 下载源码与权重参数

去GitHub下载YOLOv5源码压缩包和相关的权重文件，我下载的是预训练好的yolov5s.pt文件。
将yolov5s.pt放于yolov5根目录下。
用pycharm打开yolov5项目。

3.2 在pycharm里面添加配置好的环境

不过添加的时候，我这边有个bug是没有搭建好的python.exe，后来查询得知：先点击conda.bat，然后在下拉列表中选择已经搭建好的环境。

3.3 下载安装依赖

在pycharm终端输入

pip install -r requirements.txt
1

可能安装依赖会遇到很多bug，但是我这很顺利没有出现任何bug，此处就不做赘述

3.4 运行detect.py

此处我有遇到之前CPU跑项目的bug：Bad git executable错误
解决方法：
Import os后面添加一行os.environ[“GIT_PYTHON_REFRESH] = “quiet”

运行成功后，可以看到detect后的两张图，即可开始训练自己的数据。

4.训练自己的数据

此处参考 她叫谢雨路 的文章：在CPU上跑yolov5（详细步骤+适合入门）

4.1 数据集准备

自己收集数据集，使用labelImg进行打标签工作，上一篇文章有讲，此处不再赘述。
我是将自己的数据集分为训练集：验证集：测试集=8：1：1，在yolov5根目录下创建own_datas文件夹，并将数据集作如下放置：

yolov5-master\data路径下找到coco128.yaml文件，复制到yolov5-master\own_datas路径下，改名为own_datas.yaml（名字可以自己随便取）
在yolov5-master\models路径下找到yolov5s.yaml文件，同样复制到yolov5-master\own_datas路径下。

打开own_datas.yaml，修改：

# path: ../datasets/coco128  # dataset root dir
train: own_datas/images/train  # train images 
val: own_datas/images/valid  # val images 
1
2
3

以及（我这边是6类）

nc: 6  # number of classes
names: ['x1','x2','x3','x4','x5','x6']  # class names
1
2

注释这一行：

# download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
1

打开yolov5s.yaml文件，修改：

nc: 6  # number of classes
1

4.2 运行train.py

打开train.py，修改：

    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='own_datas/yolov5s.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default='own_datas/own_datas.yaml', help='dataset.yaml path')
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
1
2
3
4

运行train.py
运行的时候我有遇到一个bug：Downloading https://ultralytics.com/assets/Arial.ttf to …ubuntu18.4-YOLOv5报错[01]
解决方案：可以预先从https://ultralytics.com/assets/Arial.ttf这里将Arial.ttf下载放在yolov5的根目录下，就可以避免Arial.ttf字体会自动下载。

运行：GPU是真的快，4s一个epoch，我CPU跑10min一个epoch，栓q

5.预测

训练完成后在yolov5-master\runs\train\exp文件夹中查看训练相关的信息，yolov5-master\runs\train\exp\weights文件夹里面的best.pt即为训练的最好的模型。
测试文件在yolov5-master\own_datas\images\test，打开detect.py，修改：

    parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp/weights/best.pt', help='model path(s)')
    parser.add_argument('--source', type=str, default='own_datas/test', help='file/dir/URL/glob, 0 for webcam')
1
2

结束后可以在yolov5-master\runs\detect\exp文件夹里查看结果

结束语

浅浅记录从0安装环境、安装软件、搭建环境、训练数据集、预测的一整套流程，也记录了整套流程下来遇到的一些bug。