这里接上一篇在本机电脑安装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