您的位置:首页 > 娱乐 > 明星 > tensorrt动态batch推理注意事项

tensorrt动态batch推理注意事项

2024/10/5 20:17:13 来源:https://blog.csdn.net/m0_48095841/article/details/139991828  浏览:    关键词:tensorrt动态batch推理注意事项

一、背景:使用pytorch进行训练得到pt模型, 然后使用torch.onnx把pt模型转化为onnx模型。然后再使用tensorrt自带的trtexec.exe文件把onnx模型转化为engine文件。

(1)在使用C++进行推理的时候发现一个batch的数据,值推理了里面的第一张,其他图片输出的结果都为0;

原因分析:这是因为在使用torch.onnx把pt文件转为onnx的时候,没有给他增加动态batch,导致默认的batch为1,所以在推理一个batch里面的多个数据时,只对第一个数据进行了加速推理操作,其他的数据就没进入推理,所以其他数据输出为0;

(2)解决方法:

1.首先pt模型转onnx模型的时候,在export加上 dynamic_axes 参数,这样在把onnx转为tensorrt的engine文件的时候,才能转化成功。

import torch.onnxmodel = torch.load("cls_model.pth")
model.eval()dummy_input = torch.randn(1, 3, 64, 64).to("cuda")
torch.onnx.export(model,dummy_input,"cls_model_dyn.onnx",export_params=True,input_names=["input"],output_names=["output"],verbose=True, # verbose=True 会打印出详细的转换过程dynamic_axes={'input':{0:"batch"}, "output":{0:"batch"}} # 设置动态的batch_size)

2.然后在使用trtexec.exe把onnx转成engine文件的话,给它设置动态范围大小minShapes、optShapes、maxShapes、(根据自己推理大小来更改,我的是64)

 .\trtexec.exe --onnx=cls_model_dyn.onnx --saveEngine=cls.engine --explicitBatch --minShapes=input:1x3x64x64 --optShapes=input:8x3x64x64 --maxShapes=input:8x3x64x64

这样子在使用C++推理的时候就可以推理一个batch的数据了。

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com