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)