Windows Python PyTorch CUDA 11.7 TensorRT 环境配置

博文目录

---

版本说明

截止到 2022.12.24, 相关工具情况如下

• Nvidia GeForce Game Ready 驱动程序: 527.56, 运行 nvidia-smi 可知该驱动最高已支持到最新的 CUDA 12
• Nvidia CUDA: 最新版 CUDA 版本为 12
• Nvidia TensorRT: TensorRT 8.5 GA Update 1, 支持 CUDA 11.0 到 11.8
• Nvidia cuDNN: cuDNN 8.7.0, for CUDA 11.x
• PyTorch: Windows 上支持 Python 3.7 到 3.10. 最高支持 CUDA 11.7

版本选择

pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com

目前我们无法从 Nvidia 官方 Python Index 库安装 python-tensorrt, 好在有替代方法

Nvidia TensorRT 内置了 tensorrt-8.5.2.2-cp310-none-win_amd64.whl, 可以直接安装到 Python 虚拟环境中, 支持 Python 3.6 到 3.10

综上所述, 我们选择 Python 3.10 和 CUDA 11.7 来搭建环境

下载代码

Yolo v5 7.0

下载或克隆代码到某个路径, 如 C:\mrathena\develop\workspace\pycharm\yolov5-7.0

创建并激活虚拟环境

Windows Python PyCharm 开发环境搭建

使用 Conda 创建并管理虚拟环境

conda create -n gpu python=3.10 # 创建环境
conda activate gpu # 激活环境

conda remove -n gpu --all # 删除环境

使用 CPU 推理

安装工程运行的最少依赖

cd C:\mrathena\develop\workspace\pycharm\yolov5-7.0 # 切换工作路径到工程下
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple # 安装运行依赖

运行 detect.py

出现如 YOLOv5 2022-11-22 Python-3.10.8 torch-1.13.1+cpu CPU 的字样即说明环境配置成功

首次运行会自动下载权重文件 yolov5s.pt, 也可自行下载并放在工程根目录下

执行结果见 runs\detect\exp

使用 Nvidia GPU 推理

默认依赖安装好后, Yolo 即可以以 CPU 的方式运行, 若想以 GPU 的方式运行, 还需配置 CUDA 环境

安装 PyTorch CUDA 环境

在 PyTorch 拿到 pytorch 的安装命令

conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

原版命令如果卡在某一步, 无法成功安装, 可尝试使用下方的替代方法

# 配置清华源
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
# 修改后的命令, 明确指定了各个模块的版本
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia
# 参考内容见 https://blog.csdn.net/cxin917/article/details/127825180
# 版本关系见 https://github.com/pytorch/vision 和 https://github.com/pytorch/audio, 和执行 `conda search pytorch-cuda`

安装完成后, 执行 python, 输入如下内容, 如果返回 True, 则说明 CUDA 环境配置成功. 执行 exit() 退出 python 命令

import torch
torch.cuda.is_available()

运行 detect.py

有一个报错, 应该是 pillow 的原因

PIL（Python Image Library），Python 图像处理标准库。仅支持到 Python 2.7，故 Python 3 用的是 PIL 的兼容版本 Pillow

C:\mrathena\develop\miniconda\envs\gpu\lib\site-packages\torch\storage.py:11: UserWarning: The NumPy module was reloaded (imported a second time). This can in some cases result in small but subtle issues and is discouraged.
  import numpy as np
Traceback (most recent call last):
  File "C:\mrathena\develop\workspace\pycharm\yolov5-7.0\detect.py", line 43, in <module>
    from models.common import DetectMultiBackend
  File "C:\mrathena\develop\workspace\pycharm\yolov5-7.0\models\common.py", line 25, in <module>
    from PIL import Image
  File "C:\mrathena\develop\miniconda\envs\gpu\lib\site-packages\PIL\Image.py", line 100, in <module>
    from . import _imaging as core
ImportError: DLL load failed while importing _imaging: 找不到指定的模块。

无需卸载, 通过如下命令直接强制替换问题版本即可. 切记不要使用 conda uninstall pillow 这个命令, 因为会删除其他很多东西

pip install pillow==9.2.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

出现如 YOLOv5 2022-11-22 Python-3.10.8 torch-1.13.1 CUDA:0 (NVIDIA GeForce RTX 2080, 8192MiB) 的字样即说明环境配置成功

安装 TensorRT

百度网盘 相关资源

# 本地安装 tensorrt 模块, 从百度网盘中下载对应 Python 版本的安装包即可
pip install tensorrt-8.5.2.2-cp310-none-win_amd64.whl
# 安装 onnx, 需要先将 pt 转为 onnx 再转为 engine
pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

导出 engine

执行下方命令将 yolov5s.pt 导出为 yolov5s.engine

python export.py --weights yolov5s.pt --device 0 --include engine

运行 detect.py

将 detect.py 中的 weights 参数的默认值 yolov5s.pt 修改为 yolov5s.engine, 然后再运行

C:\mrathena\develop\miniconda\envs\gpu\python.exe C:/mrathena/develop/workspace/pycharm/yolov5-7.0/detect.py
detect: weights=yolov5s.engine, source=data\images, data=data\coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs\detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5  2022-11-22 Python-3.10.8 torch-1.13.1 CUDA:0 (NVIDIA GeForce RTX 2080, 8192MiB)

Loading yolov5s.engine for TensorRT inference...
[12/24/2022-22:30:04] [TRT] [I] [MemUsageChange] Init CUDA: CPU +410, GPU +0, now: CPU 14174, GPU 1213 (MiB)
[12/24/2022-22:30:04] [TRT] [I] Loaded engine size: 31 MiB
[12/24/2022-22:30:05] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +699, GPU +258, now: CPU 14961, GPU 1503 (MiB)
[12/24/2022-22:30:05] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +0, now: CPU 0, GPU 0 (MiB)
[12/24/2022-22:30:05] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 14930, GPU 1503 (MiB)
[12/24/2022-22:30:05] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +0, now: CPU 0, GPU 0 (MiB)
image 1/2 C:\mrathena\develop\workspace\pycharm\yolov5-7.0\data\images\bus.jpg: 640x640 4 persons, 1 bus, 4.6ms
image 2/2 C:\mrathena\develop\workspace\pycharm\yolov5-7.0\data\images\zidane.jpg: 640x640 2 persons, 2 ties, 5.0ms
Speed: 720.5ms pre-process, 4.8ms inference, 3.9ms NMS per image at shape (1, 3, 640, 640)
Results saved to runs\detect\exp6

额外配置

百度网盘 相关资源

报缺少 nvinfer.dll / nvinfer_plugin.dll / nvonnxparser.dll / nvparsers.dll / cudnn64_8.dll / cublas64_11.dll / cublasLt64_11.dll 等做如下操作

下载百度网盘中的 lib 文件夹到某个位置, 并将该位置添加到 Path 环境变量

或自行下载 Nvidia TensorRT, 将其中的 lib 文件夹解压到某个位置, 并将该位置添加到 Path 环境变量, 下载百度网盘中 lib 文件夹中的 cudnn64_8.dll, 放到该 lib 文件夹中, 保证 Path 能覆盖到