RetinaNet是一種One-Stage RCNN類型的檢測網絡,基本結構由一個Backbone、多個子網及NMS后處理組成。許多訓練框架中均實現了RetinaNet,典型的框架有Detectron2。本文以Detectron2的標準RetinaNet實現為例,介紹如何使用Blade優化RetinaNet(Detectron2)類型的模型。

使用限制

本文使用的環境需要滿足以下版本要求:
  • 系統環境:Linux系統中使用Python 3.6及其以上版本、CUDA 10.2。
  • 框架:PyTorch 1.8.1及其以上版本、Detectron2 0.4.1及其以上版本。
  • 推理優化工具:Blade 3.16.0及其以上版本。

操作流程

使用Blade優化RetinaNet(Detectron2)類型模型的流程如下:
  1. 步驟一:導出模型

    使用Detectron2提供的TracingAdapterscripting_with_instances任何一種方式導出模型。

  2. 步驟二:調用Blade優化模型

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

  3. 步驟三:加載運行優化后的模型

    經過對優化前后的模型進行性能測試,如果對結果滿意,可以加載優化后的模型進行推理。

步驟一:導出模型

Detectron2FAIR開源的靈活、可擴展、可配置的目標檢測和圖像分割訓練框架。由于框架的靈活性,使用常規方法導出模型可能會失敗或得到錯誤的導出結果。為了支持TorchScript部署,Detectron2提供了TracingAdapterscripting_with_instances兩種導出方式,詳情請參見Detectron2 Usage

Blade支持輸入任意形式的TorchScript模型,如下以scripting_with_instances為例,介紹導出模型的過程。
import torch
import numpy as np

from torch import Tensor
from torch.testing import assert_allclose

from detectron2 import model_zoo
from detectron2.export import scripting_with_instances
from detectron2.structures import Boxes
from detectron2.data.detection_utils import read_image

# 使用scripting_with_instances導出RetinaNet模型。
def load_retinanet(config_path):
    model = model_zoo.get(config_path, trained=True).eval()
    fields = {
        "pred_boxes": Boxes,
        "scores": Tensor,
        "pred_classes": Tensor,
    }
    script_model = scripting_with_instances(model, fields)
    return model, script_model

# 下載一張示例圖片。
# 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)))

# 嘗試執行和對比導出模型前后的結果。
pytorch_model, script_model = load_retinanet("COCO-Detection/retinanet_R_50_FPN_3x.yaml")
with torch.no_grad():
    batched_inputs = [{"image": img.float()}]
    pred1 = pytorch_model(batched_inputs)
    pred2 = script_model(batched_inputs)

assert_allclose(pred1[0]['instances'].scores, pred2[0].scores)

步驟二:調用Blade優化模型

  1. 調用Blade優化接口。
    調用blade.optimize接口對模型進行優化,代碼示例如下。關于blade.optimize接口詳情,請參見優化PyTorch模型
    import blade
    
    test_data = [(batched_inputs,)] # PyTorch的輸入數據是List of tuple。
    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: (3, 480, 640) data type: float32"
      },
      "optimizations": [
        {
          "name": "PtTrtPassFp16",
          "status": "effective",
          "speedup": "3.77",
          "pre_run": "40.64 ms",
          "post_run": "10.78 ms"
        }
      ],
      "overall": {
        "baseline": "40.73 ms",
        "optimized": "10.76 ms",
        "speedup": "3.79"
      },
      "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(pytorch_model, batched_inputs)
    # 對優化后的模型測速。
    benchmark(optimized_model, batched_inputs)
    本次測試的參考結果值如下。
    Latency: 42.38
    Latency: 10.77
    上述結果表示同樣執行200輪,優化前后的模型平均延時分別是42.38 ms10.77 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 detectron2
    import torch
    
    from torch.testing import assert_allclose
    from detectron2.utils.testing import (
        get_sample_coco_image,
    )
    
    pytorch_model = model_zoo.get("COCO-Detection/retinanet_R_50_FPN_3x.yaml", trained=True).eval()
    optimized_model = torch.jit.load('optimized.pt')
    
    img = read_image('./input.jpg')
    img = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))
    
    with torch.no_grad():
        batched_inputs = [{"image": img.float()}]
        pred1 = pytorch_model(batched_inputs)
        pred2 = optimized_model(batched_inputs)
    
    assert_allclose(pred1[0]['instances'].scores, pred2[0].scores, rtol=1e-3, atol=1e-2)