onnx转trt方法

第一种:使用TensorRT自带的trtexec.exe程序自动转换

第一步:进入TensorRT的bin文件里面

cd C:\TensorRT\TensorRT-8.5.2.2\bin

第二步:输入命令 执行trtexec.exe程序

trtexec.exe --onnx=models.onnx --saveEngine=models.engine --fp16

其中models.onnx是需要转换的onnx文件,models.engine是转换完成后TensorRT可以使用的引擎文件.fp16 是精度,这里选择fp16是因为在有速度的前提下保证精度不会变化太大(精度变化为0.001左右)

第二种:自己写程序程序转换

程序源码如下:

# 导入必用依赖
import tensorrt as trt
#onnx文件路径设置
onnx_path="M:\deep\deeplabv3-plus\deeplabv3-plus-pytorch-main\model_data\models2.onnx"
#引擎文件保存路径设置
engine_path=r"M:\deep\deeplabv3-plus\tensorRT_pridect\models_2.engine"
# 创建logger：日志记录器
logger = trt.Logger(trt.Logger.WARNING)

# 创建构建器builder
builder = trt.Builder(logger)
# 预创建网络
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# 加载onnx解析器
parser = trt.OnnxParser(network, logger)
success = parser.parse_from_file(onnx_path)
for idx in range(parser.num_errors):
    print(parser.get_error(idx))
if not success:
    pass  # Error handling code here
# builder配置
config = builder.create_builder_config()
# 分配显存作为工作区间，一般建议为显存一半的大小
config.max_workspace_size = 12 << 30  # 1 Mi
serialized_engine = builder.build_serialized_network(network, config)
# 序列化生成engine文件
with open(engine_path, "wb") as f:
    f.write(serialized_engine)
    print("generate file success!")