AI 助理
备案控制台
文档
输入文档关键字查找
首页 人工智能平台 PAI 实践教程 AI加速 推理加速（Blade） RetinaNet优化案例3：结合Blade和TensorRT Plugin优化模型

RetinaNet优化案例3：结合Blade和TensorRT Plugin优化模型

更新时间：
产品详情
相关技术圈
我的收藏

大部分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后处理的decode和nms的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的模型。

    以下代码中通过RetinaNetWrapper、RetinaNetBackboneAndHeads及RetinaNetPostProcess重新组合了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)
上一篇：RetinaNet优化案例2：结合Blade和Custom C++ Operator优化模型下一篇：BERT模型优化案例：使用Blade优化基于TensorFlow的BERT模型
文档推荐