0%

TensorRT 教程(三):PyTorch模型转ONNX模型

此系列为 PyTorch model 转 TRT engine 系列第三章。

至于为什么选 PyTorch 而不是 TensorFlow,是因为笔者对 PyTorch 最为熟悉,另外 PyTorch 的易用性和动态图特点,使得在学术界也广泛采用,新的模型更新也 release 较快。本文所用的开源项目包含:detectron2ONNXONNX-simplifierTensorRT

1. 安装 ONNX

关于ONNX的介绍,请参见其官网。简单来说,ONNX 是一种充作中间转换的角色,将训练时候的不同开源库训练得到的model 转换成为不同的另一种 model 的表示方法。最常见的就是本文即将提到的 PyTorch 转 ONNX 再转 TensorRT 的工作方式。ONNX 本身还带有可运行的底层库 ONNX runtime,不过一般不使用。

考虑到我们后面对 ONNX 需要做一些自定义操作,因此建议源码安装,方便进行修改。在 detectron2 文件夹下新建文件夹 thirdlibs,进入文件夹运行如下命令:

  1. 下载 ONNX: git clone https://github.com/onnx/onnx.git
  2. 更新子模块: cd onnx && git submodule update --init --recursive
  3. 安装: python setup.py develop # 由于频繁更改,所以建议使用 develop,不使用 install
  4. 测试: python -c "import onnx"

如果不报错即安装成功。

2. ResNet50 的 PyTorch 模型

使用 detectron2 训练 Faster R-CNN 教程可参考官网教程,此处我们使用 detectron2 MODEL_ZOO 中的已经训练好的 Faster R-CNN 模型,其Backbone 即为 ResNet50,config 文件为 detectron2 默认配置:../configs/PascalVOC-Detection/faster_rcnn_R_50_FPN.yaml。关于 Faster R-CNN 的网络结构可参看本站前一篇博文:Mask R-CNN 模型结构详解 。当然,我们也可以先不考虑数值的正确性,即我们只随机初始化模型之后并不训练网络,这样可以快速的测试 ONNX 和 TRT 对 op 的支持情况。通常在训练自己的数据期间,进行模型部署的开发,提高效率。

3. PyTorch 模型转 ONNX 模型

Detectron2 在早先版本并不支持 Faster R-CNN 的转换,而且只支持转到 caffe2 模型。Caffe2 模型的转换也是通过 ONNX 作为媒介进行,这里使用 detectron2 转到 Caffe2 的中间ONNX模型,是因为 Caffe2 的中间转换做了很多 ONNX 的 model 图优化,对某些 op 进行了 merge。当然也可不使用该方式进行,使用 ONNX 源生的 op ,但是转换得到的 ONNX 模型可视化出来会有 op 没有进行优化。

3.1 利用 Detectron2 自带 Caffe2 模型转换模块

detectron2/projects 路径下新建文件夹命名为 fasterrcnn_trt,然后新建 python 文件convert_fasterrcnn_onnx.py,拷贝detectron2/tools/deploy/caffe2_converter.py 内容到该文件中,并对函数稍作修改如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
...

if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Convert a model to Caffe2")
parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
parser.add_argument("--run-eval", action="store_true")
parser.add_argument("--output", help="output directory for the converted caffe2 model")
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
logger = setup_logger()
logger.info("Command line arguments: " + str(args))

cfg = setup_cfg(args)

# create a torch model
torch_model = build_model(cfg)
DetectionCheckpointer(torch_model).resume_or_load(cfg.MODEL.WEIGHTS)


...

faster_rcnn_trt 文件夹中运行如下命令:

1
python convert_fasterrcnn_onnx.py --config-file ../../configs/PascalVOC-Detection/faster_rcnn_R_50_FPN.yaml --output ./onnx_model --run-eval MODEL.WEIGHTS detectron2://COCO-Detection/faster_rcnn_R_50_FPN_1x/137257794/model_final_b275ba.pkl MODEL.DEVICE cpu

首先会从 MODEL_ZOO 中下载已经训练好的 Faster R-CNN 模型。然后进行 parse ,会看到输出的 log 信息,大致如下:

1
2
3
4
5
6
7
8
9
10
......
%318 : Float(1, 3, 800, 1216) = onnx::Cast[to=1](%316)
%319 : Float(1, 3, 800, 1216) = onnx::Sub(%318, %_wrapped_model.pixel_mean)
%320 : Float(1, 3, 800, 1216) = onnx::Div(%319, %_wrapped_model.pixel_std)
%input.1 : Float(1, 3, 800, 1216) = _caffe2::AliasWithName[name="normalized_data", is_backward=0](%320)
%322 : Float(1, 64, 400, 608) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[7, 7], pads=[3, 3, 3, 3], strides=[2, 2]](%input.1, %_wrapped_model.backbone.bottom_up.stem.conv1.weight)
%323 : Float(1, 64, 400, 608) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%322, %_wrapped_model.backbone.bottom_up.stem.conv1.norm.weight, %_wrapped_model.backbone.bottom_up.stem.conv1.norm.bias, %_wrapped_model.backbone.bottom_up.stem.conv1.norm.running_mean, %_wrapped_model.backbone.bottom_up.stem.conv1.norm.running_var)
%324 : Float(1, 64, 400, 608) = onnx::Relu(%323)
%325 : Float(1, 64, 200, 304) = onnx::MaxPool[kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[2, 2]](%324)
......

我们可以用 netron 对生成的 ONNX 模型进行可视化:

3.2 使用源生 ONNX 进行转换

为了达到训练过程的友好,detectron2 在forward时对输入数据进行了封装,以其中一张图为例:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
[{'file_name': 'coco/val2017/000000000139.jpg', 'height': 426, 'width': 640, 'image_id': 139, 'image': tensor([[[ 73,  74,  76,  ...,  39,  38,  37],
[ 74, 75, 77, ..., 40, 39, 38],
[ 76, 77, 78, ..., 41, 40, 39],
...,
[152, 152, 153, ..., 109, 98, 92],
[151, 151, 152, ..., 76, 67, 62],
[150, 150, 151, ..., 57, 49, 45]],

[[136, 138, 141, ..., 75, 72, 71],
[137, 139, 142, ..., 76, 73, 72],
[140, 141, 143, ..., 78, 75, 74],
...,
[183, 183, 184, ..., 105, 89, 80],
[183, 183, 184, ..., 72, 62, 55],
[183, 183, 184, ..., 54, 46, 41]],

[[170, 171, 172, ..., 68, 69, 70],
[171, 172, 173, ..., 69, 70, 70],
[172, 173, 174, ..., 71, 71, 71],
...,
[186, 186, 187, ..., 181, 176, 173],
[186, 186, 187, ..., 144, 136, 131],
[186, 186, 187, ..., 123, 113, 107]]], dtype=torch.uint8)}]

如果直接使用上面输出导出 ONNX 模型,会遇到如下错误:

1
2
...
RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type int

因此,需要对 detectron2 进行定制化修改,使其适应 ONNX 模型的导出要求。

需要做的工作:

  1. 修改 detectron2 build 模型时的接口,使 forward过程输入不需要支持输入 image 的 filename等值
  2. 修改 input ,使 input 为列表或字典,包括以下信息: feature map,image size,缩放比例信息

以 backbone ResNet50 为例:

首先新建 class :

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class ResNet50(torch.nn.Module):

def __init__(self, cfg, torch_model):
selfcfg = cfg
self.model = torch_model
self.eval()

def forward(self, input):
images = input["images"]
features = self.model.backbone(images)
return features

def convert_input(self, inputs):
images = self.model.preprocess_image(batched_inputs)
min_size = 800
max_size = 1333
for i in range(len(batched_inputs)):
s = max(batched_inputs[i]["height"] * 1.0 / min_size,
batched_inputs[i]["width"] * 1.0 / max_size)
scales.append((s,))
im_info = []
for image_size, scale in zip(images.image_sizes, scales):
im_info.append([*image_size, *scale])
return {
"images": images.tensor,
"image_sizes": torch.tensor(images.image_sizes),
"im_info": torch.tensor(im_info, device=images.tensor.device),
}

经过以上修改,即可导出 ONNX 源生 model。或者更改 ONNX 默认输入方式,使其支持 detectron2 的输入类型。

Conclusion

本文介绍了 ONNX 和在 detectron2 框架下导出 Faster R-CNN ONNX 模型的两种方式。

Reference

  1. ONNX
  2. PyTorch
  3. Python发布工具setuptools的用法
  4. detectron2 MODEL_ZOO