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)類型模型的流程如下:
- 步驟一:導出模型
使用Detectron2提供的
TracingAdapter
或scripting_with_instances
任何一種方式導出模型。 - 步驟二:調用Blade優化模型
調用
blade.optimize
接口優化模型,并保存優化后的模型。 - 步驟三:加載運行優化后的模型
經過對優化前后的模型進行性能測試,如果對結果滿意,可以加載優化后的模型進行推理。
步驟一:導出模型
Detectron2是FAIR開源的靈活、可擴展、可配置的目標檢測和圖像分割訓練框架。由于框架的靈活性,使用常規方法導出模型可能會失敗或得到錯誤的導出結果。為了支持TorchScript部署,Detectron2提供了TracingAdapter
和scripting_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)