Platform For AI:Optimize RetinaNet with Blade and custom C++ operators

1. 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

2. Compile the custom C++ operators

PyTorch supports three build methods for custom operators: CMake, Just-in-Time (JIT) compilation, and Setuptools. The following table summarizes when to use each:

This tutorial uses JIT compilation. The following code compiles decode.cu and nms.cu into custom.so inside the retinanet-examples folder:

import torch.utils.cpp_extension
import os

codebase = "retinanet-examples"
sources = [
    'csrc/extensions.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="custom",
    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_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,
)

After this runs, custom.so is saved in retinanet-examples/.

3. Replace the Python post-processing with custom operators

The two compiled operators are:

• torch.ops.retinanet.decode — decodes anchor boxes and class scores

• torch.ops.retinanet.nms — applies NMS and returns final detections

The adapter_forward function below uses these operators in place of RetinaNet's original Python post-processing. It is then patched onto RetinaNet.forward:

import os
import torch
from typing import Tuple, Dict, List, Optional

codebase = "retinanet-examples"
torch.ops.load_library(os.path.join(codebase, 'custom.so'))

decode_cuda = torch.ops.retinanet.decode
nms_cuda = torch.ops.retinanet.nms

# Same as RetinaNet.forward, but uses decode_cuda and nms_cuda for post-processing.
def adapter_forward(self, batched_inputs: Tuple[Dict[str, torch.Tensor]]):
    images = self.preprocess_image(batched_inputs)
    features = self.backbone(images.tensor)
    features = [features[f] for f in self.head_in_features]
    cls_heads, box_heads = self.head(features)
    cls_heads = [cls.sigmoid() for cls in cls_heads]
    box_heads = [b.contiguous() for b in box_heads]

    # Decode anchor boxes and scores per feature level.
    strides = [images.tensor.shape[-1] // cls_head.shape[-1] for cls_head in cls_heads]
    decoded = [
        decode_cuda(
            cls_head,
            box_head,
            anchor.view(-1),
            stride,
            self.test_score_thresh,
            self.test_topk_candidates,
        )
        for stride, cls_head, box_head, anchor in zip(
            strides, cls_heads, box_heads, self.cell_anchors
        )
    ]

    # Concatenate across feature levels, then apply NMS.
    decoded = [torch.cat(tensors, 1) for tensors in zip(decoded[0], decoded[1], decoded[2])]
    return nms_cuda(decoded[0], decoded[1], decoded[2], self.test_nms_thresh, self.max_detections_per_image)

from detectron2.modeling.meta_arch import retinanet

# Patch the model class with the custom-operator-based forward function.
retinanet.RetinaNet.forward = adapter_forward

1. Run the optimization

Call blade.optimize on the exported TorchScript model. Load the custom operator library first — Blade needs to trace through it during optimization.

import os
import blade
import torch

# Load the custom operator library before optimization.
codebase = "retinanet-examples"
torch.ops.load_library(os.path.join(codebase, 'custom.so'))

blade_config = blade.Config()
blade_config.gpu_config.disable_fp16_accuracy_check = True

test_data = [(batched_inputs,)]  # PyTorch input is a list of tuples.

with blade_config:
    optimized_model, opt_spec, report = blade.optimize(
        script_model,       # TorchScript model from Step 2.
        'o1',               # O1 optimization level.
        device_type='gpu',  # Target device.
        test_data=test_data,
    )

For a full description of blade.optimize parameters, see Optimize a PyTorch model.

2. Print the report and save the model

# Print the optimization report.
print("Report: {}".format(report))

# Save both models for later comparison.
torch.jit.save(script_model, 'script_model.pt')
torch.jit.save(optimized_model, 'optimized.pt')

The optimization report shows which passes were applied and their impact:

{
  "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: (3, 480, 640) data type: float32"
  },
  "optimizations": [
    {
      "name": "PtTrtPassFp16",
      "status": "effective",
      "speedup": "3.92",
      "pre_run": "40.72 ms",
      "post_run": "10.39 ms"
    }
  ],
  "overall": {
    "baseline": "40.64 ms",
    "optimized": "10.41 ms",
    "speedup": "3.90"
  },
  "model_info": {"input_format": "torch_script"},
  "compatibility_list": [{"device_type": "gpu", "microarchitecture": "T4"}],
  "model_sdk": {}
}

The PtTrtPassFp16 optimization — TensorRT FP16 conversion — produces the ~3.9x speedup. For a full field reference, see Optimization report.

3. Run a performance test

The following benchmark function runs 100 warmup iterations followed by 200 timed iterations:

import time

@torch.no_grad()
def benchmark(model, inp):
    # Warmup
    for i in range(100):
        model(inp)
    torch.cuda.synchronize()
    # Timed run
    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, batched_inputs)    # Before optimization
benchmark(optimized_model, batched_inputs) # After optimization

Expected output:

Latency: 40.65
Latency: 10.46

1. Authenticate

# Optional: prevent unexpected exits during the trial period.
export BLADE_AUTH_USE_COUNTING=1

# Required: set your region and authentication token.
export BLADE_REGION=<region>
export BLADE_TOKEN=<token>

Replace the placeholders:

Obtain both values by joining the PAI-Blade DingTalk user group. For the QR code, see Install PAI-Blade.

2. Load and verify

import blade.runtime.torch
import detectron2
import torch
import numpy as np
import os
from detectron2.data.detection_utils import read_image
from torch.testing import assert_allclose

# Load the custom operator library before loading the model.
codebase = "retinanet-examples"
torch.ops.load_library(os.path.join(codebase, 'custom.so'))

script_model = torch.jit.load('script_model.pt')
optimized_model = torch.jit.load('optimized.pt')

img = read_image('./input.jpg')
img = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))

# Verify that outputs are numerically equivalent.
with torch.no_grad():
    batched_inputs = [{"image": img.float()}]
    pred1 = script_model(batched_inputs)
    pred2 = optimized_model(batched_inputs)

assert_allclose(pred1[0], pred2[0], rtol=1e-3, atol=1e-2)

The tolerance values (rtol=1e-3, atol=1e-2) account for rounding differences introduced by FP16 computation.