基于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年開源的目標檢測網絡?；赥ransformer的目標檢測算法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

六后續計劃

●適配更多基于Transformer網絡的目標檢測、分割任務算法模型；

●適配視覺大模型（SAM、DINOv2）。

審核編輯：湯梓紅