Windows Python PyTorch CUDA 11.7 TensorRT 环境配置

mrathena

已于 2024-06-05 19:47:07 修改

阅读量9k

点赞数 11

分类专栏： python 文章标签： python windows pytorch

于 2022-12-24 20:08:22 首次发布

本文链接：https://blog.csdn.net/mrathena/article/details/128430943

版权

python 专栏收录该内容

7 篇文章 40 订阅

订阅专栏

博文目录

文章目录

版本说明
版本选择
下载代码
创建并激活虚拟环境
使用 CPU 推理
- 安装工程运行的最少依赖
- 运行 detect.py
使用 Nvidia GPU 推理
安装 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\yolo v5 7.0

创建并激活虚拟环境

Windows Python PyCharm 开发环境搭建

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

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

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

使用 CPU 推理

安装工程运行的最少依赖

cd C:\mrathena\develop\workspace\pycharm\yolo v5 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 的安装命令

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

运行 detect.py

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

20240605 补充

如果报错 NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend, 原因是自动安装的 torchvision 版本不对, 应该是自动安装了 cpu 版本, 执行 pip list 查看 torch 组件版本, 如果错误, 需要手动替换 torchvision 为 gpu 版本, 如 torchvision-0.18.0+cu121-cp311-cp311-win_amd64.whl, 下载地址在这里

torch                   2.3.0+cu121
torchaudio              2.3.0+cu121
# 错误的版本
torchvision             0.18.0
# 正确的版本
torchvision             0.18.0+cu121

安装 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

20240605 补充

如果 onnx 和 protobuf 版本冲突, 将 protobuf 的版本指定去除即可

导出 engine

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

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

运行 detect.py

将 detect.py 中的 weights 参数的默认值 yolov5s.pt 修改为 yolov5s.engine, 然后再运行, 耗时只有原先的 1/10 左右

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

额外配置

百度网盘相关资源

以 TensorRT 方式运行, 需要一些其他依赖, 可通过下载对应 zip 包, 从中找到对应运行库, 并添加到环境变量 Path 中

Nvidia TensorRT: nvinfer.dll / nvinfer_plugin.dll / nvonnxparser.dll / nvparsers.dll
Nvidia cuDNN: cudnn64_8.dll
Nvidia CUDA: cublas64_11.dll / cublasLt64_11.dll

也可以从我提供的百度网盘中下载对应的文件