基于AX650N部署DETR

一背景

目标检测作为计算机视觉应用的基础任务，一直备受行业应用重点关注。从最早学术界的Two-Stage方案Faster-RCNN，到首个利于工业界高效率部署的One-Stage方案SSD（Single Shot MultiBox Detector），最后发展到这两年大家非常熟悉的YOLO系列（v1/v2/v3/v4/v5/v6/v7/v8/x……），这一系列经典检测器方案主要是基于卷积神经网络（CNN）作为特征提取的Backbone，然后使用手工组件Anchor-Base（Faster-RCNN、SSD、YOLOv1-v7）或者Anchor-Free（YOLOv8、YOLOX）加上非最大抑制（NMS）来筛选出最终目标框。然而Anchor-Base或Anchor-Free的两种方案都利用非最大抑制进行后处理，这给经典检测器带来了推理性能的瓶颈。此外，由于非最大抑制不使用图像信息，因此在边界框保留和删除中有各式NMS原理带来的问题。

近年来Transformer被广泛应用到计算机视觉的物体分类领域，例如ViT、SwinT等。那是否也能进一步将Transformer拓展到目标检测任务呢？本文通过分享Transformer用在目标检测领域的开山之作：DETR(DEtection TRansformer)，同时尝试在AX650N上完成端到端的部署，给行业内对边缘侧/端侧部署Transformer模型的爱好者提供一种新的思路。

二DETR

DETR是Meta在2020年开源的目标检测网络。基于Transformer的目标检测算法DETR，采用简洁的pipeline，去除NMS、Anchor设计，且在COCO数据集上的指标与Faster RCNN相当。

DETR的网络结构如下图所示，主要是由四个模块组成：基于ResNet的骨干网络、编码器、解码器、预测头。

 

2.1 骨干网络

采用经典的卷积神经网络ResNet50或ResNet101作为其Backbone，最终输出降采样32倍的Feature Map给Transformer的Encoder（编码器）作为输入。

 

2.2 编码器

Encoder得到backbone输入的Feature Map之后，先用过一个Conv1x1 + Reshape操作，将特征值进行序列化（sequence）变化，然后紧跟一系列的MHA和FFN处理（Transformer网络核心部件）。

 

2.3 解码器

有两个输入分别是：

● 编码器得到的特征；

● Object Queries：类似基于CNN目标检测算法中的Anchor Boxes。

 

2.4 预测头

采用FFN操作，分别输出BBox和Class ID。

三模型转换

 

3.1 Pulsar2

Pulsar2是新一代AI工具链，包含模型转换、离线量化、模型编译、异构调度四合一超强功能，进一步强化了网络模型高效部署的需求。在针对第三代NPU架构进行了深度定制优化的同时，也扩展了算子&模型支持的能力及范围，对Transformer结构的网络也有较好的支持。

pulsar2 deploy pipeline

 

3.2 模型下载

我们从DTER的官方仓库获取对应的ONNX模型。

下载github仓库

git仓库下载

 

git clone https://github.com/facebookresearch/detr.git
cd detr
pip install -r requirements.txt

 

修改main.py文件

修改main.py

 

$ git diff main.py
diff --git a/main.py b/main.py
index e5f9eff..bf7855c 100644
--- a/main.py
+++ b/main.py
@@ -139,7 +139,7 @@ def main(args):
                                   weight_decay=args.weight_decay)
     lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)


-    dataset_train = build_dataset(image_set='train', args=args)
+    dataset_train = build_dataset(image_set='val', args=args)
     dataset_val = build_dataset(image_set='val', args=args)


     if args.distributed:
@@ -180,6 +180,9 @@ def main(args):
             optimizer.load_state_dict(checkpoint['optimizer'])
             lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
             args.start_epoch = checkpoint['epoch'] + 1
+
+    torch.onnx.export(model,torch.ones(1,3,608,608).to(args.device),f="detr_r50_608.onnx",do_constant_folding=True,opset_version=11)
+    exit()


     if args.eval:
         test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,

 

导出ONNX模型

导出ONNX模型

 

$ python main.py --batch_size 2 --no_aux_loss --eval --resume https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth --coco_path /home/xigua/data/coco2017 --device cpu
$ onnxsim detr_r50_608.onnx detr_r50_608-sim.onnx

 

 

3.3 模型编译

一键完成图优化、离线量化、编译、对分功能。整个过程耗时不到5分钟，相比上一代工具链模型编译效率有了数量级的提升。

编译log

 

$ pulsar2 build --input model/detr_r50_608-sim.onnx --output_dir output --config config/detr_config.json
patool: Extracting ./dataset/dataset_v04.zip ...
patool: ... ./dataset/dataset_v04.zip extracted to `output/quant/dataset/samples'.
                                                                                   Quant Config Table
┏━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Input   ┃ Shape            ┃ Dataset Directory ┃ Data Format ┃ Tensor Format ┃ Mean                                              ┃ Std                                               ┃
┡━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ samples │ [1, 3, 608, 608] │ samples           │ Image       │ BGR           │ [123.68000030517578, 116.77899932861328,          │ [58.619998931884766, 57.34000015258789,           │
│         │                  │                   │             │               │ 103.93900299072266]                               │ 57.599998474121094]                               │
└─────────┴──────────────────┴───────────────────┴─────────────┴───────────────┴───────────────────────────────────────────────────┴───────────────────────────────────────────────────┘
Transformer optimize level: 1
1 File(s) Loaded.
[18:23:26] AX Refine Operation Config Pass Running ...    Finished.
[18:23:26] AX Transformer Optimize Pass Running ...       Finished.
[18:23:26] AX Quantization Config Refine Pass Running ... Finished.
[18:23:27] AX Quantization Fusion Pass Running ...        Finished.
[18:23:27] AX Quantization Simplify Pass Running ...      Finished.
[18:23:27] AX Parameter Quantization Pass Running ...     Finished.
Calibration Progress(Phase 1): 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:01<00:00,  1.86s/it]
Calibration Progress(Phase 2): 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:06<00:00,  6.14s/it]
Finished.
[1853] AX Passive Parameter Quantization Running ...  Finished.
[1853] AX Parameter Baking Pass Running ...           Finished.
[1853] AX Refine Int Parameter pass Running ...       Finished.
Network Quantization Finished.
quant.axmodel export success: output/quant/quant_axmodel.onnx
Building native ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 000
......
2023-05-10 1837.093 | INFO     | yasched.test_onepass1428 - max_cycle = 45853543
2023-05-10 18:26:23.472 | INFO     | yamain.command.build889 - fuse 1 subgraph(s)

 

 

3.4 Graph Optimize

这里的Graph Optimize主要是对转入的ONNX模型进行一系列预置好的图优化功能，利于后续编译阶段提升执行效率。

Original ONNX to Optimize ONNX

 

3.5 Graph Quantize

Pulsar2采用大家熟悉的PTQ策略对计算图进行8bit量化压缩，满足NPU硬件单元的计算需求，量化后的模型同样采用ONNX容器进行存储，方便用户使用Netron打开进行观察量化后的计算图。下面截取DETR中核心单元MHA（Mulit-Head Attention）和FFN（Feed Forward Network）。

MHA and FFN

四

上板部署

 

4.1 AX-Samples

开源项目AX-Samples实现了基于爱芯元智AI SoC的常见深度学习开源算法部署示例，方便社区开发者进行快速评估和适配。

最新版本已开始提供AX650系列的NPU示例，其中也包含了本文介绍的DETR参考代码。

https://github.com/AXERA-TECH/ax-samples/blob/main/examples/ax650/ax_detr_steps.cc

 

4.2 运行

运行log

 

root@AXERA:/home/test# ./ax_detr -m detr_r50_npu1.axmodel -i ssd_car.jpg -r 10
--------------------------------------
model file : detr_r50_npu1.axmodel
image file : ssd_car.jpg
img_h, img_w : 608 608
--------------------------------------
Engine creating handle is done.
Engine creating context is done.
Engine get io info is done.
Engine alloc io is done.
Engine push input is done.
--------------------------------------
prob_pred_idx=0 ,bbox_pred_idx=1
post process cost time:0.47 ms
--------------------------------------
Repeat 10 times, avg time 49.50 ms, max_time 49.51 ms, min_time 49.49 ms
--------------------------------------
detection num: 3
 3:  99%, [ 318,  193,  499,  325], car
 6:  96%, [ 118,   53,  455,  284], bus
 1:  93%, [ 189,  168,  273,  374], person
--------------------------------------

 

运行结果

 

五性能统计

单位：FPS

Model Name	3.6Tops@Int8	10.8Tops@Int8
DETR_R50	20	43

六后续计划

● 适配更多基于Transformer网络的目标检测、分割任务算法模型；

● 适配视觉大模型（SAM、DINOv2）。

 

　　审核编辑：汤梓红