Platform For AI:Optimize RetinaNet with Blade and TensorRT plugin

When to use this approach

The ONNX export path — exporting a PyTorch model to ONNX and then deploying with TensorRT — has significant limitations for detection models. ONNX opset support in TensorRT is limited, and the NMS post-processing is especially difficult to export reliably. Many teams implement NMS as a custom TensorRT plugin to work around this.

If you have already built a custom TensorRT plugin for post-processing, Blade can work with it directly. This lets you optimize the full model — backbone, heads, and post-processing — in a single pass, without rewriting the plugin.

If you have not yet implemented a TensorRT plugin, two simpler alternatives are available:

• Use Blade directly (case 1): Export the model with scripting_with_instances and run blade.optimize. See RetinaNet optimization case 1.

• Use Blade with a TorchScript custom C++ operator (case 2): Implement NMS as a TorchScript custom C++ operator. See RetinaNet optimization case 2.

Download the sample code

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

Compile the TensorRT plugin

The sample code includes TensorRT plugin implementations for RetinaNetDecode and RetinaNetNMS. Three compilation methods are available: CMake, Just-in-Time (JIT), and Setuptools. See Extending TorchScript with custom C++ operators for a comparison.

This tutorial uses JIT compilation.

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,
)

Encapsulate the backbone and heads

Wrap the RetinaNet convolutional layers in a separate RetinaNetBackboneAndHeads module, isolating them from the TensorRT post-processing subgraph.

import torch
from typing import List
from torch import Tensor
from torch.testing import assert_allclose
from detectron2 import model_zoo

# Encapsulates the backbone and RPN heads of RetinaNet.
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)

Build the TensorRT post-processing engine

Skip this step if you have already created a TensorRT engine.

The following code dynamically loads plugin.so and uses the TensorRT Python API to build a decode-and-NMS network, then compiles it into an engine.

import os
import numpy as np
import tensorrt as trt

import ctypes
# Load the TensorRT plugin shared library.
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

# Retrieve a TensorRT plugin by name.
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

# Build the TensorRT decode-and-NMS network.
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 from cls and box heads.
    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 the decode subgraph (one RetinaNetDecode plugin per feature level).
    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 the NMS subgraph.
    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

Create the cuda_engine based on the actual output shapes of RetinaNetBackboneAndHeads:

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)

Assemble the full model with Blade extensions

Blade's TRTEngineExtension interface wraps the compiled engine as a TorchScript-compatible module. RetinaNetWrapper combines RetinaNetBackboneAndHeads and the TRT post-processing module into a single exportable graph.

import blade.torch

# Post-processing module backed by the Blade TensorRT extension.
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)

# Full RetinaNet model: PyTorch backbone/heads + TensorRT post-processing.
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)

# Export and save as 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)

The assembled module supports TorchScript export and serialization through torch.classes.torch_addons.TRTEngineExtension.

Run blade.optimize

Call blade.optimize on the exported TorchScript model. For full parameter details, see Optimize a PyTorch model.

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,)]  # Input data must be a list of tuples.
with blade_config:
    optimized_model, opt_spec, report = blade.optimize(
        script_model,        # TorchScript model from the previous step.
        'o1',                # Blade O1 optimization level.
        device_type='gpu',   # Target device.
        test_data=test_data, # Test data for optimization and accuracy checking.
    )

Save the optimized model

The optimized model is a TorchScript model. Print the optimization report and save it:

# Print the optimization report.
print("Report: {}".format(report))
# Save the optimized model.
torch.jit.save(optimized_model, 'optimized.pt')

The report output looks like this. For field descriptions, see Optimization 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": {}
}

Benchmark the original and optimized models

Run the benchmark to compare latency:

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 the original model.
benchmark(script_model, example_inputs)
# Benchmark the optimized model.
benchmark(optimized_model, example_inputs)

Reference results (200 runs, NVIDIA T4):

Latency: 40.71
Latency: 9.35

Average latency drops from 40.71 ms to 9.35 ms — a 4.32x speedup.

Authenticate Blade

1. (Trial only) Set the following environment variable to prevent unexpected exits caused by authentication failures:

   export BLADE_AUTH_USE_COUNTING=1

2. Set your region and token:

   export BLADE_REGION=<region>
   export BLADE_TOKEN=<token>

   To get the <region> and <token> values, join the PAI-Blade DingTalk user group. The QR code is available in Install PAI-Blade.

Run inference

The optimized model is a standard TorchScript model. Load and run it without switching environments:

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)