YOLOV5 GPU训练_yolov5使用gpu训练命令

1. 查看是否支持GPU
   • 打开cmd，输入以下指令：

   nvidia-smi
   1

   

图中的红框的CUDA版本代表的是Cuda driver version，是与我们显卡的驱动相关的
2. 安装或者查看自己的显卡驱动

https://www.nvidia.cn/Download/index.aspx?lang=cn
1

最好是将自己的显卡驱动更新到最新版本

3. 显卡的算力必须与Cuda runtime version相匹配

https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#id4
1

https://en.wikipedia.org/wiki/CUDA#cite_note-38
1

3060ti对应的算力是8.6
再查看显卡算力8.6对应的Cuda runtime version

对应的版本是11.1-11.4

4. 下载pytorch

https://pytorch.org/
1

这里只有11.7和11.8两个版本，不符合我的要求，点击图中的install previous versions去寻找合适的版本。

运行以下命令换源安装pytorch

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c http://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/win-64
1

进入python环境后通过以下命令验证环境：

import torch
print(torch.backends.cudnn.enabled)
1
2

返回true表示成功
验证以下Cuda

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#返回cuda表示成功
#或者
print(torch.cuda.is_available())
#返回True表示成功
1
2
3
4
5
6

5. git yolov5代码 后，进入到requirements.txt文件所在目录，用过以下命令安装所需库：

pip install -r requirements.txt
1

到此，环境配置结束。
使用GPU训练模型的命令：

python train.py --data mask_data.yaml --cfg mask_yolov5s.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 4 --device 0
1

将pt文件转为onnx

python export.py --weights ./pretrained/yolov5s.pt
1