最近一段时间一直在负责转模型为tensorRT, 鉴于我们部门的模型现在基本上都是pytorch的了,其他框架暂时不考虑。这里总结一 下经验与方法,并放一个使用network API 转换 resnet18的例子,这个例子是我做另一个网络模型转换时的其中的骨干网络,其他部分差不太多,不过因为是工作中的内容,就不放出来了,这里使用resnet18为例。
对一个pytorch 模型,通常可以使用官方的方法,把它转成onnx,并使用TensorRT官方的,trtexec 工具转成模型,或使用网络定义api重写网络并赋值权重,先说几个经验:
- 如果可以使用官方的trtexec 直接把你的从pytorch导出onnx模型转成trt,那这是最方便快捷的,所以先测试下这个路行不行(注意opset版本要与你使用的trt兼容)
- 如果trtexec不能直接转onnx. (例如pytorch里,上采样算子 upsample 与torch.nn.functional.interpolate,对于LINEAR模式的,有align_corners属性的常常有问题)可以考虑使用 onnx2trt ,我在我一个模型上尝试过,但精度在模型加载完我自己的权重后与原版本的pytorch不一样,不加载是一样的。总之奇怪的问题比较多时,考虑使用网络定义。
调试方法
你环境正常的情况下,是可以像普通程序一样调试的,当然构建过程可能看不了,不过trt的报错很清楚,只要注意给每层命名,你就能知道是哪个层出的问题,然后对照原代码,检查kernelsize 什么的参数是不是设置错了,一般来说如果设置错会报错说形状不一样的。
关键点:
- 注意给每层命名
- 如果哪层有问题,从后向前面注释掉些层,检查哪里的问题
- 也可以逐层检查输出结果,很简单,你把一个中间层设置为输出结点,对原代码与现在生成的都输入同样的数据,查看结果是不是一样(对pytorch的debug就可以看到一些中间结果了)。 比如都输入一个tensorshape为(1,3,256,256)全为1,两边shape 一定要一样,不然他们的输出结果不一样很正常。
不知道咋实现的层:
找找官方api吧(https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Network.html#) ,其实应该都有,比如我想找两个tensor相加,就没注意到是add_elementwise,点开看看IElementWiseLayer
的说明,就发现是它了,它可以对两个tensor做各种算术运算。
ResNet18的转换
先看下pytorch 的resnet18的核心代码,具体代码可以去我的github上查看,请点:RESNET18。
pytorch 的resnet18 核心代码:
class ResNet(nn.Module):
def __init__(self, block, layers, in_channels=3, dcn=None):
self.dcn = dcn
self.inplanes = 64
super(ResNet, self).__init__()
self.out_channels = []
self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(bl