Windows Python PyTorch CUDA 11.7 TensorRT 环境配置

时间:2022-12-25 07:55:47

博文目录


版本说明

截止到 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 能覆盖到