yolo 搭建GPU训练环境

安装流程

1、查看cuda版本
2、安装cuda Toolkit
3、安装cuDNN
4、安装pytorch
5、安装yolo其他依赖
6、运行测试

查看当前显卡的Cuda版本

C:\workspace\yolov5>nvidia-smi
Tue Nov 21 17:21:35 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 526.47       Driver Version: 526.47       CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA RTX A4000   WDDM  | 00000000:01:00.0 Off |                  Off |
| 41%   47C    P8    20W / 140W |   1586MiB / 16376MiB |      8%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

环境下载

CUDA版本是12.0，向下兼容，CUDA Toolkit下载11.8，cuDNN下载8.8.1.3

安装pytorch
conda install pytorch2.1.0 torchvision0.16.0 torchaudio==2.1.0 pytorch-cuda=11.8 -c pytorch -c nvidia

安装yolo依赖
pip3 install -r requirements.txt

软件下载
CUDA Toolkit
https://developer.nvidia.com/cuda-toolkit-archive

cuDNN
https://developer.nvidia.com/rdp/cudnn-archive

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

安装指定版本的torch

pip install torch2.1.0+cu121 -i https://download.pytorch.org/whl/cu121
pip install torchaudio2.1.0+cu121 -i https://download.pytorch.org/whl/cu121
pip install torchvision==0.16.0+cu121 -i https://download.pytorch.org/whl/cu121

验证cuda环境

import torch

print("是否可用：", torch.cuda.is_available())  # 查看GPU是否可用
print("GPU数量：", torch.cuda.device_count())  # 查看GPU数量
print("torch方法查看CUDA版本：", torch.version.cuda)  # torch方法查看CUDA版本
print("GPU索引号：", torch.cuda.current_device())  # 查看GPU索引号
print("GPU名称：", torch.cuda.get_device_name(0))  # 根据索引号得到GPU名称