基于AX650N部署SegFormer

发布时间:2024-09-02  

背景

语义分割(Semantic Segmentation)是计算机视觉中的一项基本任务。与单张图像分类任务相比,语义分割相当于是像素级别上的分类任务。语义分割为许多下游应用特别是近几年来的智能驾驶技术的落地提供了可能。

本文将简单介绍SegFormer的基本原理,同时指导如何导出ONNX模型,并将其部署在优秀的端侧AI芯片AX650N上,希望能给行业内对边缘侧/端侧部署Transformer模型的爱好者提供新的思路。


SegFormer介绍

SegFormer的论文中提出了一个简单、高效的语义分割方案,它将Transformers与轻量级多层感知(MLPs)解码器结合起来。SegFormer有两个吸引人的特点:

1

SegFormer包含一个新的层次结构的Transformer编码器,输出多尺度特征。它不需要位置编码,这样就不用对位置编码做插值。

2

SegFormer避免了复杂的解码器。所提出的MLP解码器从不同的层聚集信息,从而结合Local Attention和Global Attention来呈现强大的表示。这种简单而轻量级的设计是用Transformer高效地做分割的关键。

论文中扩展了上述方案,得到了一系列不同大小型号的模型,从SegFormer-B0到SegFormer-B5,相比之前的分割模型达到了更好的性能和效率。例如,SegFormer-B4在ADE20K上以64M参数实现了50.3%的mIoU,最佳模型SegFormer-B5在Cityscapes验证集上达到了84.0%的mIoU。

2.1 骨干网络

5dedd428-0f45-11ee-962d-dac502259ad0.png

backbone

2.2 分级Transformer编码器

论文中提出了一系列的Mix Transformer编码器(MiT),MiT-B0到MiT-B5,具有相同的结构,但尺寸不同。MiT-B0是用于快速推理的轻量级模型,而MiT-B5是用于最佳性能的最大模型。设计的MiT部分灵感来自ViT,但针对语义分割进行了定制和优化。

2.3 轻量级All-MLP解码器

集成了一个仅由MLP层组成的轻量级解码器,这避免了其他方法中通常用的手工制作和计算要求很高的组件。实现这种简单解码器的关键是分级Transformer编码器比传统的CNN编码器具有更大的有效感受野(ERF)。

5e3be17c-0f45-11ee-962d-dac502259ad0.png

benchmark


AX650N

AX650N是一款兼具高算力与高能效比的SoC芯片,集成了八核Cortex-A55 CPU,10.8TOPs@INT8 NPU(针对 Transformer 模型进行了定制优化),支持8K@30fps的ISP,以及H.264、H.265编解码的VPU。接口方面,AX650N支持64bit LPDDR4x,多路MIPI输入,千兆Ethernet、USB、以及HDMI 2.0b输出,并支持32路1080p@30fps解码。强大的性能可以让AX650N帮助用户在智慧城市、智慧教育、智能制造等领域发挥更大的价值。

模型转换

本文以segformer-b0-cityscapes-640-1280为例。

4.1 模型下载

这次我们推荐从Huggingface的ModelZoo下载模型。

5ea3bd4c-0f45-11ee-962d-dac502259ad0.png

huggingface

● ONNX模型导出的脚本

ONNX模型导出


import torch

from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor

from pathlib import Path

from onnxruntime.quantization import quantize_dynamic, QuantType, preprocess

import onnx

import onnxruntime

import os

from PIL import Image

from typing import List



def export_model(model_name: str, export_dir: str, input_sample: torch.Tensor):

  model = SegformerForSemanticSegmentation.from_pretrained(model_name)

  model.eval()

  export_path = os.path.join(export_dir, model_name)

  Path(export_path).mkdir(parents=True, exist_ok=True)

  onnx_path = os.path.join(export_path, 'model.onnx')

  # export the model to ONNX while preserving the first dimension as dynamic

  torch.onnx.export(model, input_sample, onnx_path, export_params=True, 

            opset_version=11,

            input_names=['input'],

            output_names=['output'],

            )



export_dir = './hf_export/'



model_name = 'nvidia/segformer-b0-finetuned-cityscapes-640-1280'

export_model(model_name, export_dir, torch.randn([1,3,640,1280]))



# model_name = 'nvidia/segformer-b1-finetuned-ade-512-512'

# export_model(model_name, export_dir, torch.randn([1,3,512,512]))

# check_models(model_name, export_dir, image_paths)

● onnxsim优化


$ onnxsim segformer-b0-cityscapes-640-1280.onnx segformer-b0-cityscapes-640-1280-sim.onnx

Simplifying...

Finish! Here is the difference:

┏━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓

┃      ┃ Original Model ┃ Simplified Model ┃

┡━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩

│ Add    │ 136      │ 136       │

│ Concat   │ 21       │ 1        │

│ Constant  │ 176      │ 0        │

│ Conv    │ 20       │ 20        │

│ Div    │ 46       │ 46        │

│ Erf    │ 8       │ 8        │

│ MatMul   │ 68       │ 68        │

│ Mul    │ 46       │ 46        │

│ Pow    │ 30       │ 30        │

│ ReduceMean │ 60       │ 60        │

│ Relu    │ 1       │ 1        │

│ Reshape  │ 76       │ 76        │

│ Resize   │ 4       │ 4        │

│ Shape   │ 20       │ 0        │

│ Slice   │ 20       │ 0        │

│ Softmax  │ 8       │ 8        │

│ Sqrt    │ 30       │ 30        │

│ Sub    │ 30       │ 30        │

│ Transpose │ 76       │ 76        │

│ Model Size │ 14.3MiB    │ 14.3MiB     │

└────────────┴────────────────┴──────────────────┘

● 添加argmax输出头


由于AX650N的NPU支持argmax算子,因此可以将argmax添加到该模型的输出头,直接获取每个像素点置信度最高的类别ID。


首先安装onnx_graphsurgeon依赖:


pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

运行下面的脚本,添加argmax op:


import numpy as np

import onnx

import onnx_graphsurgeon as gs



model_path = './segformer-b0-cityscapes-640-1280-sim.onnx'

output_model_path = './segformer-b0-cityscapes-640-1280-sim-argmax.onnx'



onnx_model = onnx.load(model_path)

onnx_graph = gs.import_onnx(onnx_model)

node_last_conv = onnx_graph.nodes[-1]



# Attrs for ArgMax

axis = 1

keepdims = 1



argmax_out_shape = node_last_conv.outputs[0].shape.copy()

argmax_out_shape[axis] = 1



argmax_out = gs.Variable(

  'argmax_output',

  dtype=np.int64,

  shape=argmax_out_shape,

)



argmax_node = gs.Node(

  op='ArgMax',

  name='decode_head_ArgMax',

  inputs=[node_last_conv.outputs[0]],

  outputs=[argmax_out],

  attrs={'axis': axis, 'keepdims': keepdims},

)



onnx_graph.nodes.append(argmax_node)

onnx_graph.outputs.clear()

onnx_graph.outputs = [argmax_out]

onnx_graph.cleanup().toposort()

onnx_model_with_argmax = gs.export_onnx(onnx_graph)

onnx_model_with_argmax.ir_version = onnx_model.ir_version



onnx.save(onnx_model_with_argmax, output_model_path)


添加argmax前后两个ONNX模型对比:

5ee9c300-0f45-11ee-962d-dac502259ad0.png

add-argmax


4.2 模型编译


使用AX650N配套的AI工具链Pulsar2,一键完成图优化、离线量化、编译、对分功能。


$ pulsar2 build --input model/segformer-b0-cityscapes-640-1280-sim-argmax.onnx --output_dir segformer/ --config config/seg          former_config.json --npu_mode NPU3

Building onnx ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00

patool: Extracting ./dataset/coco_4.tar ...

patool: running /usr/bin/tar --extract --file ./dataset/coco_4.tar --directory segformer/quant/dataset/input

patool: ... ./dataset/coco_4.tar extracted to `segformer/quant/dataset/input'.

                                       Quant Config Table

┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓

┃ Input ┃ Shape       ┃ Dataset Directory ┃ Data Format ┃ Tensor Format ┃ Mean                     ┃ Std                     ┃

┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩

│ input │ [1, 3, 640, 1280] │ input       │ Image    │ RGB      │ [123.67500305175781, 116.27999877929688,   │ [58.39500045776367, 57.119998931884766,   │

│    │          │          │       │        │ 103.52999877929688]              │ 57.356998443603516]             │

└───────┴───────────────────┴───────────────────┴─────────────┴───────────────┴───────────────────────────────────────────────┴──────────────────────────────────────────────┘

4 File(s) Loaded.

[14:17:44] AX LSTM Operation Format Pass Running ...   Finished.

[14:17:44] AX Refine Operation Config Pass Running ...  Finished.

[14:17:44] AX Transformer Optimize Pass Running ...    Finished.

[14:17:45] AX Reset Mul Config Pass Running ...      Finished.

[14:17:45] AX Tanh Operation Format Pass Running ...   Finished.

[14:17:45] AX Quantization Config Refine Pass Running ... Finished.

[14:17:45] AX Quantization Fusion Pass Running ...    Finished.

[14:17:45] AX Quantization Simplify Pass Running ...   Finished.

[14:17:45] AX Parameter Quantization Pass Running ...   Finished.

Calibration Progress(Phase 1): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:23<00:00,  5.86s/it]

Finished.

[1411] AX Passive Parameter Quantization Running ...  Finished.

[1411] AX Parameter Baking Pass Running ...           Finished.

[1412] AX Refine Int Parameter Pass Running ...       Finished.

Network Quantization Finished.

quant.axmodel export success: segformer/quant/quant_axmodel.onnx

Building native ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 000

......

2023-05-30 1418.661 | INFO     | yasched.test_onepass1615 - max_cycle = 37883954

2023-05-30 14:35:04.545 | INFO     | yamain.command.build904 - fuse 1 subgraph(s)


上板部署


5.1 AX-Samples


开源项目AX-Samples实现了常见的深度学习开源算法在爱芯元智的AI SoC上的示例代码,方便社区开发者进行快速评估和适配。最新版本已开始提供AX650系列的NPU示例,其中也包含了本文介绍的SegFormer参考代码。


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


5.2 运行


# ./ax_segformer -m segformer-b0-cityscapes-640-1280-argmax.axmodel -i segformer_test.png

--------------------------------------

model file : segformer-b0-cityscapes-640-1280-argmax.axmodel

image file : segformer_test.png

img_h, img_w : 640 1280

--------------------------------------

post process cost time:7.07 ms

--------------------------------------

Repeat 1 times, avg time 48.15 ms, max_time 48.15 ms, min_time 48.15 ms

--------------------------------------

文章来源于:电子工程世界    原文链接
本站所有转载文章系出于传递更多信息之目的,且明确注明来源,不希望被转载的媒体或个人可与我们联系,我们将立即进行删除处理。

我们与500+贴片厂合作,完美满足客户的定制需求。为品牌提供定制化的推广方案、专属产品特色页,多渠道推广,SEM/SEO精准营销以及与公众号的联合推广...详细>>

利用葫芦芯平台的卓越技术服务和新产品推广能力,原厂代理能轻松打入消费物联网(IOT)、信息与通信(ICT)、汽车及新能源汽车、工业自动化及工业物联网、装备及功率电子...详细>>

充分利用其强大的电子元器件采购流量,创新性地为这些物料提供了一个全新的窗口。我们的高效数字营销技术,不仅可以助你轻松识别与连接到需求方,更能够极大地提高“闲置物料”的处理能力,通过葫芦芯平台...详细>>

我们的目标很明确:构建一个全方位的半导体产业生态系统。成为一家全球领先的半导体互联网生态公司。目前,我们已成功打造了智能汽车、智能家居、大健康医疗、机器人和材料等五大生态领域。更为重要的是...详细>>

我们深知加工与定制类服务商的价值和重要性,因此,我们倾力为您提供最顶尖的营销资源。在我们的平台上,您可以直接接触到100万的研发工程师和采购工程师,以及10万的活跃客户群体...详细>>

凭借我们强大的专业流量和尖端的互联网数字营销技术,我们承诺为原厂提供免费的产品资料推广服务。无论是最新的资讯、技术动态还是创新产品,都可以通过我们的平台迅速传达给目标客户...详细>>

我们不止于将线索转化为潜在客户。葫芦芯平台致力于形成业务闭环,从引流、宣传到最终销售,全程跟进,确保每一个potential lead都得到妥善处理,从而大幅提高转化率。不仅如此...详细>>