gpu训练yolo模型

1,安装cuda

查看nvidia版本(控制面板 -> nvcuda64.dll)

https://developer.nvidia.com/cuda-toolkit-archive

验证

cmd -> nvcc --version

2,安装cudann

https://developer.nvidia.com/rdp/cudnn-download

注册

下载

解压

将三个文件夹拷贝到到cuda的安装目录下

添加环境变量

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\include

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

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\libnvvp

验证

cmd -> 进入到 extras\demo_suite下，执行.\bandwidthTest.exe和.\deviceQuery.exe

3,安装gpu版的torch

https://pytorch.org/get-started/locally

选择配置后复制--texra-index-url的路径

下载对应的torch和torchvision版本的.whl文件

pip install xxx.whl文件到项目

检查是否可以使用cuda

torch.cuda.is_available()

常见问题

内存不足

减小batch

页面太小,调大虚拟内存

window系统路径中不能出现中文,linux可以