TensorRT10.6 python版本简单使用流程

原创

已于 2025-07-21 11:44:50 修改 · 1.2k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#python #人工智能 #深度学习 #pytorch

于 2024-11-28 11:28:55 首次发布

该文章已生成可运行项目，

TensorRT python版本 windows上下载安装：
TensorRT python Windows 下载安装

写在前面：

经过阅读多篇文章、查看官方文档、论坛，终于成功打通了python版本的TensorRT流程。
本文介绍的是python版本的tensorRT推理流程，测试pytroch框架保存的分类模型。本人综合测试推理速度比pytorch快一倍以上，该数据仅供参考。

python==3.9.20 onnx==1.17.0 numpy==1.24.4 tensorrt==10.6.0

整体流程：

头文件

import tensorrt as trt
import torch
import pycuda.driver as cuda
import pycuda.autoinit
from mobilenetv3 import MobileNetV3_Large  #分类网络
import numpy as np
import cv2

1.根据pytorch模型生成onnx模型
需要安装onnx库。

#推荐使用加载模型参数的方法。
model = torch.load('model.pth')
# 切换到评估模式，禁用 Dropout 等层
model.eval()  
# 根据模型的输入形状调整
dummy_input = torch.randn(1, 3, 112, 112).to(device)
onnx_model_path = "model.onnx"
torch.onnx.export(model, dummy_input, onnx_model_path,
                  input_names=['input'], output_names=['output'],
                  opset_version=11