基于AX650N部署DETR

一背景

目標(biāo)檢測(cè)作為計(jì)算機(jī)視覺(jué)應(yīng)用的基礎(chǔ)任務(wù)，一直備受行業(yè)應(yīng)用重點(diǎn)關(guān)注。從最早學(xué)術(shù)界的Two-Stage方案Faster-RCNN，到首個(gè)利于工業(yè)界高效率部署的One-Stage方案SSD（Single Shot MultiBox Detector），最后發(fā)展到這兩年大家非常熟悉的YOLO系列（v1/v2/v3/v4/v5/v6/v7/v8/x……），這一系列經(jīng)典檢測(cè)器方案主要是基于卷積神經(jīng)網(wǎng)絡(luò)（CNN）作為特征提取的Backbone，然后使用手工組件Anchor-Base（Faster-RCNN、SSD、YOLOv1-v7）或者Anchor-Free（YOLOv8、YOLOX）加上非最大抑制（NMS）來(lái)篩選出最終目標(biāo)框。然而Anchor-Base或Anchor-Free的兩種方案都利用非最大抑制進(jìn)行后處理，這給經(jīng)典檢測(cè)器帶來(lái)了推理性能的瓶頸。此外，由于非最大抑制不使用圖像信息，因此在邊界框保留和刪除中有各式NMS原理帶來(lái)的問(wèn)題。

近年來(lái)Transformer被廣泛應(yīng)用到計(jì)算機(jī)視覺(jué)的物體分類領(lǐng)域，例如ViT、SwinT等。那是否也能進(jìn)一步將Transformer拓展到目標(biāo)檢測(cè)任務(wù)呢？本文通過(guò)分享Transformer用在目標(biāo)檢測(cè)領(lǐng)域的開(kāi)山之作：DETR(DEtection TRansformer)，同時(shí)嘗試在AX650N上完成端到端的部署，給行業(yè)內(nèi)對(duì)邊緣側(cè)/端側(cè)部署Transformer模型的愛(ài)好者提供一種新的思路。

二DETR

DETR是Meta在2020年開(kāi)源的目標(biāo)檢測(cè)網(wǎng)絡(luò)。基于Transformer的目標(biāo)檢測(cè)算法DETR，采用簡(jiǎn)潔的pipeline，去除NMS、Anchor設(shè)計(jì)，且在COCO數(shù)據(jù)集上的指標(biāo)與Faster RCNN相當(dāng)。

DETR的網(wǎng)絡(luò)結(jié)構(gòu)如下圖所示，主要是由四個(gè)模塊組成：基于ResNet的骨干網(wǎng)絡(luò)、編碼器、解碼器、預(yù)測(cè)頭。

2.1 骨干網(wǎng)絡(luò)

采用經(jīng)典的卷積神經(jīng)網(wǎng)絡(luò)ResNet50或ResNet101作為其Backbone，最終輸出降采樣32倍的Feature Map給Transformer的Encoder（編碼器）作為輸入。

2.2 編碼器

Encoder得到backbone輸入的Feature Map之后，先用過(guò)一個(gè)Conv1x1 + Reshape操作，將特征值進(jìn)行序列化（sequence）變化，然后緊跟一系列的MHA和FFN處理（Transformer網(wǎng)絡(luò)核心部件）。

2.3 解碼器

有兩個(gè)輸入分別是：

●編碼器得到的特征；

● Object Queries：類似基于CNN目標(biāo)檢測(cè)算法中的Anchor Boxes。

2.4 預(yù)測(cè)頭

采用FFN操作，分別輸出BBox和Class ID。

三模型轉(zhuǎn)換

3.1 Pulsar2

Pulsar2是新一代AI工具鏈，包含模型轉(zhuǎn)換、離線量化、模型編譯、異構(gòu)調(diào)度四合一超強(qiáng)功能，進(jìn)一步強(qiáng)化了網(wǎng)絡(luò)模型高效部署的需求。在針對(duì)第三代NPU架構(gòu)進(jìn)行了深度定制優(yōu)化的同時(shí)，也擴(kuò)展了算子&模型支持的能力及范圍，對(duì)Transformer結(jié)構(gòu)的網(wǎng)絡(luò)也有較好的支持。

pulsar2 deploy pipeline

3.2 模型下載

我們從DTER的官方倉(cāng)庫(kù)獲取對(duì)應(yīng)的ONNX模型。

下載github倉(cāng)庫(kù)

git倉(cāng)庫(kù)下載

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,

導(dǎo)出ONNX模型

導(dǎo)出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 模型編譯

一鍵完成圖優(yōu)化、離線量化、編譯、對(duì)分功能。整個(gè)過(guò)程耗時(shí)不到5分鐘，相比上一代工具鏈模型編譯效率有了數(shù)量級(jí)的提升。

編譯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主要是對(duì)轉(zhuǎn)入的ONNX模型進(jìn)行一系列預(yù)置好的圖優(yōu)化功能，利于后續(xù)編譯階段提升執(zhí)行效率。

Original ONNX to Optimize ONNX

3.5 Graph Quantize

Pulsar2采用大家熟悉的PTQ策略對(duì)計(jì)算圖進(jìn)行8bit量化壓縮，滿足NPU硬件單元的計(jì)算需求，量化后的模型同樣采用ONNX容器進(jìn)行存儲(chǔ)，方便用戶使用Netron打開(kāi)進(jìn)行觀察量化后的計(jì)算圖。下面截取DETR中核心單元MHA（Mulit-Head Attention）和FFN（Feed Forward Network）。

MHA and FFN

四

上板部署

4.1 AX-Samples

開(kāi)源項(xiàng)目AX-Samples實(shí)現(xiàn)了基于愛(ài)芯元智AI SoC的常見(jiàn)深度學(xué)習(xí)開(kāi)源算法部署示例，方便社區(qū)開(kāi)發(fā)者進(jìn)行快速評(píng)估和適配。

最新版本已開(kāi)始提供AX650系列的NPU示例，其中也包含了本文介紹的DETR參考代碼。

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

4.2 運(yùn)行

運(yùn)行l(wèi)og

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
--------------------------------------

運(yùn)行結(jié)果

五性能統(tǒng)計(jì)

單位：FPS

六后續(xù)計(jì)劃

●適配更多基于Transformer網(wǎng)絡(luò)的目標(biāo)檢測(cè)、分割任務(wù)算法模型；

●適配視覺(jué)大模型（SAM、DINOv2）。

審核編輯：湯梓紅