ONNX模型在不同显卡占用内存不同的问题:
问题描述:MobileNetv2在1660S上占用1G,在2080Ti上占用2G,在3080上高达2.6G,显存占用异常。
原因:在py文件中导入了torchvision.transforms包,导致了内存占用,删去改用自己手写normalization之后,显存占用正常了
TensorRT转换错误的BUG
1. all concat input tensors must have the same dimensions except on the concatenation axis(x)
或者
Assertion failed:ctx -->tensors().count(inputName)
报错信息:all concat input tensors must have the same dimensions except on the concatenation axis(x)
报错原因:TensorRT还没实现F.interpolate中mode为bilinear且align_corners为false的情况
解决办法:要从ONNX转TensorRT做量化,如果模型中存在F.interpolate,那么就不能用mode=“bilinear”和align_corners=True. 设置为bilinear必须把align_corners设为False. 但这样效果很差,有几种替代方法:
- nn.PixelShuffle代替F.interpolate
- nn.ConvTranspose2d代替F.interpolate
- 手写F.interpolate
由于前两种方法太麻烦,且需要重新训练,因此我采用了第三种。代码为:
x_resized = x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
expand(x.size(0), x.size(1), x.size(2), x_l.size()[2] // x.size(2), x.size(3), x_l.size()[3] // x.size(3)).\
contiguous().view(x.size(0), x.size(1), x_l.size()[2], x_l.size()[3])
除此之外,用torch.onnx.export转成onnx模型的话,op_version需要设置为10,设为11且用F.interpolate的话,在转成TensorRT的时候,会报错。
2. ONNX转TensorRT的时候申请workspace空间过小,导致解析失败
可以尽量取大一点的内存,我取了4个G才跑通。
如有错误,敬请指出,一起交流。