一、背景:使用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的数据了。