大部分PyTorch用户会使用TensorRT Plugin实现检测模型的后处理部分,以支持整个模型导出到TensorRT。Blade拥有良好的可扩展性,如果您已经自己实现了TensorRT Plugin,也可以结合Blade协同优化。本文介绍如何使用Blade对已经实现了TensorRT Plugin机制的检测模型进行优化。

背景信息

TensorRT是NVIDIA GPU平台进行推理优化的利器,Blade底层优化深度采纳了TensorRT的优化手段。相比而言,Blade有机融合了计算图优化、TensorRT/oneDNN等Vendor优化库、AI编译优化、Blade手工优化算子库、Blade混合精度及Blade EasyCompression等多种优化技术。

RetinaNet是一种One-Stage RCNN类型的检测网络,基本结构由一个Backbone、多个子网及NMS后处理组成。许多训练框架中均实现了RetinaNet,典型的框架有Detectron2。之前介绍了如何通过scripting_with_instances方式导出RetinaNet(Detectron2)模型并使用Blade快速完成模型优化,详情请参见RetinaNet优化案例1:使用Blade优化RetinaNet(Detectron2)模型

然而,对于大部分PyTorch用户而言,先导出ONNX再使用TensorRT部署是常见且熟悉的使用方式。但是ONNX导出和TensorRT对ONNX Opset的支持均有限,导致很多情况下导出ONNX并使用TensorRT优化的过程并不具备鲁棒性。特别是对于Detection网络的后处理部分,难以直接导出ONNX并使用TensorRT优化。除此之外,实际场景中检测模型的后处理部分代码实现通常不高效,因此,许多用户会使用TensorRT提供的Plugin机制实现后处理部分,以支持整个模型导出到TensorRT。

相比而言,Blade结合TorchScript Custom C++ Operators的优化方式比使用TensorRT提供的Plugin机制实现后处理部分更加简便,详情请参见RetinaNet优化案例2:结合Blade和Custom C++ Operator优化模型。此外,Blade拥有良好的可扩展性,如果您已经自己实现了TensorRT Plugin,也可以结合Blade协同优化。

使用限制

本文使用的环境需要满足以下版本限制:
  • 系统环境:Linux系统中使用Python 3.6及其以上版本、GCC 5.4及其以上版本、Nvidia Tesla T4、CUDA 10.2、CuDNN 8.0.5.39、TensorRT 7.2.2.3。
  • 框架:PyTorch 1.8.1及其以上版本、Detectron2 0.4.1及其以上版本。
  • 推理优化工具:Blade 3.16.0及其以上版本(动态链接TensorRT版本)。

操作流程

结合Blade和TensorRT Plugin优化模型的流程如下:
  1. 步骤一:创建带有TensorRT Plugin的PyTorch模型

    使用TensorRT Plugin实现RetinaNet的后处理部分。

  2. 步骤二:调用Blade优化模型

    调用blade.optimize接口优化模型,并保存优化后的模型。

  3. 步骤三:加载运行优化后的模型

    经过对优化前后的模型进行性能测试,如果对结果满意,可以加载优化后的模型进行推理。

步骤一:创建带有TensorRT Plugin的PyTorch模型

Blade能够和TensorRT扩展机制协同优化,以下介绍如何使用TensorRT扩展实现RetinaNet的后处理部分。 关于开发和编译TensorRT Plugin的教程请参见NVIDIA Deep Learning TensorRT Documentation。本文使用的RetinaNet后处理部分的程序逻辑来自NVIDIA开源社区,详情请参见Retinanet-Examples。本文抽取了核心的代码用于说明开发实现Custom Operator的流程。

  1. 下载示例代码并解压。
    wget -nv https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/tutorials/retinanet_example/retinanet-examples.tar.gz -O retinanet-examples.tar.gz
    tar xvfz retinanet-examples.tar.gz 1>/dev/null
  2. 编译TensorRT Plugin。
    示例代码中包含了RetinaNet后处理的decodenms的TensorRT Plugin实现及注册。PyTorch官方文档中(详情请参见EXTENDING TORCHSCRIPT WITH CUSTOM C++ OPERATORS)提供了三种编译Custom Operators的方式:Building with CMake、Building with JIT Compilation及Building with Setuptools。这三种编译方式适用于不同场景,您可以根据自己的需求进行选择。本文为了简便,采用Building with JIT Compilation方式,示例代码如下所示。
    说明 编译之前,您需要配置好TensorRT、CUDA,、CUDNN等依赖库。
    import torch.utils.cpp_extension
    import os
    
    codebase="retinanet-examples"
    sources=['csrc/plugins/plugin.cpp',
             'csrc/cuda/decode.cu',
             'csrc/cuda/nms.cu',]
    sources = [os.path.join(codebase,src) for src in sources]
    torch.utils.cpp_extension.load(
        name="plugin",
        sources=sources,
        build_directory=codebase,
        extra_include_paths=['/usr/local/TensorRT/include/', '/usr/local/cuda/include/', '/usr/local/cuda/include/thrust/system/cuda/detail'],
        extra_cflags=['-std=c++14', '-O2', '-Wall'],
        extra_ldflags=['-L/usr/local/TensorRT/lib/', '-lnvinfer'],
        extra_cuda_cflags=[
            '-std=c++14', '--expt-extended-lambda',
            '--use_fast_math', '-Xcompiler', '-Wall,-fno-gnu-unique',
            '-gencode=arch=compute_75,code=sm_75',],
        is_python_module=False,
        with_cuda=True,
        verbose=False,
    )
  3. 封装RetinaNet卷积模型部分。
    将RetinaNet模型部分单独封装为一个RetinaNetBackboneAndHeads Module。
    import torch
    from typing import List
    from torch import Tensor
    from torch.testing import assert_allclose
    from detectron2 import model_zoo
    
    # 这个类封装了RetinaNet的backbone和rpn heads部分。
    class RetinaNetBackboneAndHeads(torch.nn.Module):
    
        def __init__(self, model):
            super().__init__()
            self.model = model
    
        def preprocess(self, img):
            batched_inputs = [{"image": img}]
            images = self.model.preprocess_image(batched_inputs)
            return images.tensor
    
        def forward(self, images):
            features = self.model.backbone(images)
            features = [features[f] for f in self.model.head_in_features]
            cls_heads, box_heads = self.model.head(features)
            cls_heads = [cls.sigmoid() for cls in cls_heads]
            box_heads = [b.contiguous() for b in box_heads]
            return cls_heads, box_heads
    
    retinanet_model = model_zoo.get("COCO-Detection/retinanet_R_50_FPN_3x.yaml", trained=True).eval()
    retinanet_bacbone_heads = RetinaNetBackboneAndHeads(retinanet_model)
  4. 使用TensorRT Plugin构建RetinaNet后处理网络。如果您已经创建过TensorRT Engine,可以跳过此步骤。
    1. 创建TensorRT Engine。
      为了使TensorRT Plugin生效,需要实现以下功能:
      • 通过ctypes.cdll.LoadLibrary动态加载编译好的plugin.so。
      • build_retinanet_decode通过tensorrt Python API构建后处理网络并将其Build成为Engine。
      示例代码如下。
      import os
      import numpy as np
      import tensorrt as trt
      
      import ctypes
      # 加载TensorRT Plugin动态链接库。
      codebase="retinanet-examples"
      ctypes.cdll.LoadLibrary(os.path.join(codebase, 'plugin.so'))
      
      TRT_LOGGER = trt.Logger()
      trt.init_libnvinfer_plugins(TRT_LOGGER, "")
      PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
      
      # 获取TensorRT Plugin的函数。
      def get_trt_plugin(plugin_name, field_collection):
          plugin = None
          for plugin_creator in PLUGIN_CREATORS:
              if plugin_creator.name != plugin_name:
                  continue
              if plugin_name == "RetinaNetDecode":
                  plugin = plugin_creator.create_plugin(
                      name=plugin_name, field_collection=field_collection
                  )
              if plugin_name == "RetinaNetNMS":
                  plugin = plugin_creator.create_plugin(
                      name=plugin_name, field_collection=field_collection
                  )
          assert plugin is not None, "plugin not found"
          return plugin
      
      # 构建TensorRT网络的函数。
      def build_retinanet_decode(example_outputs,
              input_image_shape,
              anchors_list,
              test_score_thresh = 0.05,
              test_nms_thresh = 0.5,
              test_topk_candidates = 1000,
              max_detections_per_image = 100,
          ):
          builder = trt.Builder(TRT_LOGGER)
          EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
          network = builder.create_network(EXPLICIT_BATCH)
          config = builder.create_builder_config()
          config.max_workspace_size = 3 ** 20
      
          cls_heads, box_heads = example_outputs
          profile = builder.create_optimization_profile()
          decode_scores = []
          decode_boxes = []
          decode_class = []
      
          input_blob_names = []
          input_blob_types = []
          def _add_input(head_tensor, head_name):
              input_blob_names.append(head_name)
              input_blob_types.append("Float")
              head_shape = list(head_tensor.shape)[-3:]
              profile.set_shape(
                   head_name, [1] + head_shape, [20] + head_shape, [1000] + head_shape)
              return network.add_input(
                  name=head_name, dtype=trt.float32, shape=[-1] + head_shape
              )
      
          # Build network inputs.
          cls_head_inputs = []
          cls_head_strides = [input_image_shape[-1] // cls_head.shape[-1] for cls_head in cls_heads]
          for idx, cls_head in enumerate(cls_heads):
              cls_head_name = "cls_head" + str(idx)
              cls_head_inputs.append(_add_input(cls_head, cls_head_name))
      
          box_head_inputs = []
          for idx, box_head in enumerate(box_heads):
              box_head_name = "box_head" + str(idx)
              box_head_inputs.append(_add_input(box_head, box_head_name))
      
          output_blob_names = []
          output_blob_types = []
          # Build decode network.
          for idx, anchors in enumerate(anchors_list):
              field_coll = trt.PluginFieldCollection([
                  trt.PluginField("topk_candidates", np.array([test_topk_candidates], dtype=np.int32), trt.PluginFieldType.INT32),
                  trt.PluginField("score_thresh", np.array([test_score_thresh], dtype=np.float32), trt.PluginFieldType.FLOAT32),
                  trt.PluginField("stride", np.array([cls_head_strides[idx]], dtype=np.int32), trt.PluginFieldType.INT32),
                  trt.PluginField("num_anchors", np.array([anchors.numel()], dtype=np.int32), trt.PluginFieldType.INT32),
                  trt.PluginField("anchors", anchors.contiguous().cpu().numpy().astype(np.float32), trt.PluginFieldType.FLOAT32),]
              )
              decode_layer = network.add_plugin_v2(
                  inputs=[cls_head_inputs[idx], box_head_inputs[idx]],
                  plugin=get_trt_plugin("RetinaNetDecode", field_coll),
              )
              decode_scores.append(decode_layer.get_output(0))
              decode_boxes.append(decode_layer.get_output(1))
              decode_class.append(decode_layer.get_output(2))
      
          # Build NMS network.
          scores_layer = network.add_concatenation(decode_scores)
          boxes_layer = network.add_concatenation(decode_boxes)
          class_layer = network.add_concatenation(decode_class)
          field_coll = trt.PluginFieldCollection([
                  trt.PluginField("nms_thresh", np.array([test_nms_thresh], dtype=np.float32), trt.PluginFieldType.FLOAT32),
                  trt.PluginField("max_detections_per_image", np.array([max_detections_per_image], dtype=np.int32), trt.PluginFieldType.INT32),]
              )
          nms_layer = network.add_plugin_v2(
             inputs=[scores_layer.get_output(0), boxes_layer.get_output(0), class_layer.get_output(0)],
             plugin=get_trt_plugin("RetinaNetNMS", field_coll),
          )
          nms_layer.get_output(0).name = "scores"
          nms_layer.get_output(1).name = "boxes"
          nms_layer.get_output(2).name = "classes"
          nms_outputs = [network.mark_output(nms_layer.get_output(k)) for k in range(3)]
          config.add_optimization_profile(profile)
          cuda_engine = builder.build_engine(network, config)
          assert cuda_engine is not None
          return cuda_engine
    2. 根据RetinaNetBackboneAndHeads的实际结果输出个数,输出类型及输出Shape创建的cuda_engine
      import numpy as np
      from detectron2.data.detection_utils import read_image
      
      !wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg
      img = read_image('./input.jpg')
      img = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))
      
      example_inputs = retinanet_bacbone_heads.preprocess(img)
      example_outputs = retinanet_bacbone_heads(example_inputs)
      
      cell_anchors = [c.contiguous() for c in retinanet_model.anchor_generator.cell_anchors]
      cuda_engine = build_retinanet_decode(
                  example_outputs, example_inputs.shape, cell_anchors)
  5. 通过Blade扩展支持混合使用PyTorch和TensorRT Engine的模型。
    以下代码中通过RetinaNetWrapperRetinaNetBackboneAndHeadsRetinaNetPostProcess重新组合了Backbone、Heads及Tensorrt Plugin后处理部分。
    import blade.torch
    
    # 使用Blade TensorRT扩展支持的后处理部分。
    class RetinaNetPostProcess(torch.nn.Module):
        def __init__(self, cuda_engine):
            super().__init__()
            blob_names = [cuda_engine.get_binding_name(idx) for idx in range(cuda_engine.num_bindings)]
            input_blob_names = blob_names[:-3]
            input_blob_types = ["Float"] * len(input_blob_names)
            output_blob_names = blob_names[-3:]
            output_blob_types = ["Float"] * len(output_blob_names)
    
            self.trt_ext_plugin = torch.classes.torch_addons.TRTEngineExtension(
                bytes(cuda_engine.serialize()),
                (input_blob_names, output_blob_names, input_blob_types, output_blob_types),
            )
    
        def forward(self, inputs: List[Tensor]):
            return self.trt_ext_plugin.forward(inputs)
    
    # 混合使用PyTorch和TensorRT Engine。
    class RetinaNetWrapper(torch.nn.Module):
    
        def __init__(self, model, trt_postproc):
            super().__init__()
            self.backbone_and_heads = model
            self.trt_postproc = torch.jit.script(trt_postproc)
    
        def forward(self, images):
            cls_heads, box_heads = self.backbone_and_heads(images)
            return self.trt_postproc(cls_heads + box_heads)
    
    trt_postproc = RetinaNetPostProcess(cuda_engine)
    retinanet_mix_trt = RetinaNetWrapper(retinanet_bacbone_heads, trt_postproc)
    
    # 可以导出和保存为TorchScript。
    retinanet_script = torch.jit.trace(retinanet_mix_trt, (example_inputs, ), check_trace=False)
    torch.jit.save(retinanet_script, 'retinanet_script.pt')
    torch.save(example_inputs, 'example_inputs.pth')
    outputs = retinanet_script(example_inputs)
    新组装的torch.nn.Module拥有以下特点:
    • 使用了Blade的TensorRT扩展支持torch.classes.torch_addons.TRTEngineExtension接口。
    • 支持TorchScript模型导出,上述代码中使用了torch.jit.trace进行导出。
    • 支持TorchScript格式保存模型。

步骤二:调用Blade优化模型

  1. 调用Blade优化接口。
    调用blade.optimize接口对模型进行优化,代码示例如下。关于blade.optimize接口详情,请参见优化PyTorch模型
    import blade
    import blade.torch
    import ctypes
    import torch
    import os
    
    codebase="retinanet-examples"
    ctypes.cdll.LoadLibrary(os.path.join(codebase, 'plugin.so'))
    
    blade_config = blade.Config()
    blade_config.gpu_config.disable_fp16_accuracy_check = True
    
    script_model = torch.jit.load('retinanet_script.pt')
    example_inputs = torch.load('example_inputs.pth')
    test_data = [(example_inputs,)] # PyTorch的输入数据是List of Tuple。
    with blade_config:
        optimized_model, opt_spec, report = blade.optimize(
            script_model,  # 上一步导出的TorchScript模型。
            'o1',  # 开启Blade O1级别的优化。
            device_type='gpu',  # 目标设备为GPU。
            test_data=test_data,  # 给定一组测试数据,用于辅助优化及测试。
        )
  2. 打印优化报告并保存模型。
    Blade优化后的模型仍然是一个TorchScript模型。完成优化后,您可以通过如下代码打印优化报告并保存优化模型。
    # 打印优化结果报表。
    print("Report: {}".format(report))
    # 保存优化后的模型。
    torch.jit.save(optimized_model, 'optimized.pt')
    打印的优化报告如下所示,关于优化报告中的字段详情请参见优化报告
    Report: {
      "software_context": [
        {
          "software": "pytorch",
          "version": "1.8.1+cu102"
        },
        {
          "software": "cuda",
          "version": "10.2.0"
        }
      ],
      "hardware_context": {
        "device_type": "gpu",
        "microarchitecture": "T4"
      },
      "user_config": "",
      "diagnosis": {
        "model": "unnamed.pt",
        "test_data_source": "user provided",
        "shape_variation": "undefined",
        "message": "Unable to deduce model inputs information (data type, shape, value range, etc.)",
        "test_data_info": "0 shape: (1, 3, 480, 640) data type: float32"
      },
      "optimizations": [
        {
          "name": "PtTrtPassFp16",
          "status": "effective",
          "speedup": "4.37",
          "pre_run": "40.59 ms",
          "post_run": "9.28 ms"
        }
      ],
      "overall": {
        "baseline": "40.02 ms",
        "optimized": "9.27 ms",
        "speedup": "4.32"
      },
      "model_info": {
        "input_format": "torch_script"
      },
      "compatibility_list": [
        {
          "device_type": "gpu",
          "microarchitecture": "T4"
        }
      ],
      "model_sdk": {}
    }
  3. 对优化前后的模型进行性能测试。
    性能测试的代码示例如下所示。
    import time
    
    @torch.no_grad()
    def benchmark(model, inp):
        for i in range(100):
            model(inp)
        torch.cuda.synchronize()
        start = time.time()
        for i in range(200):
            model(inp)
        torch.cuda.synchronize()
        elapsed_ms = (time.time() - start) * 1000
        print("Latency: {:.2f}".format(elapsed_ms / 200))
    
    # 对优化前的模型测速。
    benchmark(script_model, example_inputs)
    # 对优化后的模型测速。
    benchmark(optimized_model, example_inputs)
    本次测试的参考结果值如下。
    Latency: 40.71
    Latency: 9.35
    上述结果表示同样执行200轮,优化前后的模型平均延时分别是40.71 ms和9.35 ms。

步骤三:加载运行优化后的模型

  1. 可选:在试用阶段,您可以设置如下的环境变量,防止因为鉴权失败而程序退出。
    export BLADE_AUTH_USE_COUNTING=1
  2. 获取鉴权。
    export BLADE_REGION=<region>
    export BLADE_TOKEN=<token>
    您需要根据实际情况替换以下参数:
    • <region>:Blade支持的地域,需要加入Blade用户群获取该信息,用户群的二维码详情请参见获取Token
    • <token>:鉴权Token,需要加入Blade用户群获取该信息,用户群的二维码详情请参见获取Token
  3. 加载运行优化后的模型。
    Blade优化后的模型仍然是TorchScript,因此您无需切换环境即可加载优化后的结果。
    import blade.runtime.torch
    import torch
    
    from torch.testing import assert_allclose
    import ctypes
    import os
    
    codebase="retinanet-examples"
    ctypes.cdll.LoadLibrary(os.path.join(codebase, 'plugin.so'))
    
    optimized_model = torch.jit.load('optimized.pt')
    example_inputs = torch.load('example_inputs.pth')
    
    with torch.no_grad():
        pred = optimized_model(example_inputs)