介绍

mmdetection已经成为大家进行目标检测研究与开发的首选之一，丰富的开发组件与优秀的架构使得调参与训练变成轻松有趣。

但是对于广大工程师而言，如何将算法变成实际的生产力却是个让人困扰问题。mmdetection的便利性得益于pytorch对于动态网络的强大支持，但是这却是一把双刃剑，不完全固定的图结构使得模型向其他推理框架转换变得异常困难。

TensorRT是NVIDIA提供的推理框架，支持FP16，INT8，借助dGPU和Jetson的强大算力支持，成为很多人部署模型的首选。NVIDIA官方对于pytorch模型转换TensorRT提供的方案是借助ONNX作为中间层，进行两步走的转换。对于简单的分类网络而言并没有太大问题。但是检测网络通常具有更复杂的网络结构，特殊的算子，动态的输出尺寸，使得模型转换异常困难。mmdetection初步提供了ONNX的转换，但是仅有少数模型支持，而且即使转换成为ONNX，也无法保证能够转换成TensorRT。

MMDet to TensorRT

mmdet to tensorrt就是一个为了解决上述问题而诞生的工具：

项目链接：https://github.com/grimoire/mmdetection-to-tensorrt

主要功能：

• mmdetection=>tensorrt转换：现已支持包括包括faster rcnn, cascade rcnn, ssd, retinanet等超过20种不同模型的转换，数量持续增加种。
• fp16/int8模式支持：可以轻松实现fp16模型的转换，大多数模型也已经支持int8模式
• batched input支持：以batch的形式输入数据通常可以得到更大的吞吐量，工具支持转换的模型以batch形式输入数据，只要显存吃得消，就可以享受大batch带来的速度提升。
• 动态输入尺寸：不管输入的tensor形状如何，只要mmdetection支持的，我们都支持（大概）
• 多策略组合：可以自定义自己的网络，组合各种mmdetection提供的模块，转换模型也不会是问题
• deepstream支持：转换的模型可以在deepstream中使用，已实现相关组件，享受流畅的实时流媒体目标检测吧
• 命令行与脚本支持：可以使用命令行工具轻松转换，也可以使用脚本定制你的转换。

使用方法

模型转换

CLI

mmdet2trt ${CONFIG_PATH} ${CHECKPOINT_PATH} ${OUTPUT_PATH}

PYTHON

opt_shape_param=[
    [
        [1,3,320,320],      # min shape
        [1,3,800,1344],     # optimize shape
        [1,3,1344,1344],    # max shape
    ]
]
max_workspace_size=1<<30    # some module and tactic need large workspace.
trt_model = mmdet2trt(cfg_path, weight_path, opt_shape_param=opt_shape_param, fp16_mode=True, max_workspace_size=max_workspace_size)
torch.save(trt_model.state_dict(), save_path)

使用转换后的模型进行推理：

trt_model = init_detector(save_path)
num_detections, trt_bbox, trt_score, trt_cls = inference_detector(trt_model, image_path, cfg_path, "cuda:0")

提取c++可用的序列化engine文件：

with open(engine_path, mode='wb') as f:
    f.write(model_trt.state_dict()['engine'])

模型还在不断更新中，如果有什么改进的建议欢迎提出

其他

作为项目的副产物，提供了一个可以直接转换pytorch模型的工具：
https://github.com/grimoire/torch2trt_dynamic
算是NVIDIA官方的https://github.com/NVIDIA-AI-IOT/torch2trt的威力加强版，增加了对于动态输入的支持，还有大量的新算子。本来打算PR给官方项目的，不过似乎官方为了兼容Tensorrt5暂时不打算做太大改动。因此就拿出来作为一个独立项目了。

其他2

欢迎大家PR，Issue
喜欢star，还有大家（求求你了）