这里接上一篇在本机电脑安装gpu版本的pytorch
那么安装完成之后怎么使用gpu训练自己的深度学习模型呢。

选择gpu版本的解释器

在python解释器页面找到自己安装的pytorch路径添加conda解释器并应用。

查看CUDA设备编号

在命令行输入

nvidia-smi

查看cuda设备的详细信息，如下图所示：

注意在使用gpu训练模型之前先查看python环境下是否有可用cuda设备。

import torch
torch.cuda.is_available()

如果在anaconda环境下显示true，在本机命令行下显示false，可能是没有添加环境变量。

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\lib

将lib路径放到环境变量下。
当显示True时表示cuda设备可以用了。
在训练命令后添加

--device 0

即可。

python train.py --batch-size 1 --epoch 20 --data .\data/leaf/data.yaml  --weights .\yolov5s.pt --device 0