基于AX650N部署视觉大模型DINOv2

一背景

最近一段时间，基于Transformer网络结构的视觉大模型呈现出爆发式增长，继Segment Anything（SAM）之后，Meta AI再次发布重量级开源项目——DINOv2。DINOv2可以抽取到强大的图像特征，且在下游任务上不需要微调，这使得它适合作为许多不同的应用中新的Backbone。

与之前发布的Segment Anything相比，DINOv2在应用领域和适用范围上更加广泛，原论文中的实验也涵盖了多个CV中经典的下游任务。

二DINOv2简介

在Meta AI官方的Blog中，将DINOv2的特性总结如下：

● 自监督训练

●无需Fine Tuning即可迁移至下游任务

●开箱即用的视觉基座大模型

DINOv2是一种新的高精度计算机视觉模型训练方法，使用自监督学习来实现与该领域中使用的标准方法相匹配。与其他自监督系统一样，使用DINOv2方法的模型可以在不需要任何相关原数据的情况下对任何图像集合进行训练。这意味着它可以从它所接收到的所有图像中学习，而不仅仅是那些包含特定一组标签或标题的图像。DINOv2提供了可直接用作简单线性分类器输入的高性能特征。这种灵活性意味着DINOv2可用于创建许多不同计算机视觉任务的多用途骨干。

Visualization of PCA

论文中的实验展示了DINOv2在下游任务上的出色能力，例如分类、分割和图像检索等应用领域。其中，最令人惊讶的是，在深度估计方面，DINOv2的结果明显优于in-domain与out-of-domain的SOTA的pipeline。作者认为这种强大的域外表现是自监督特征学习和轻量级任务特定模块（例如线性分类器）相结合的结果。

三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帮助用户在智慧城市、智慧教育、智能制造等领域发挥更大的价值。

四模型转换

本文以DINOv2提供的Pretrained model：ViT-S/14为例介绍部署流程。

4.1 模型下载

●下载DINOv2的开源项目

git clone https://github.com/facebookresearch/dinov2.git

●进入dinov2目录，使用下列文件替换 dinov2/models/vision_transformer.py

vision_transformer.py可从我们的开源项目中获取

https://github.com/AXERA-TECH/ax-samples/releases/download/v0.4/vision_transformer.py

●获取ONNX模型，在dinov2根目录下执行

from dinov2.models import vision_transformer as vits
import torch


_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"


def _make_dinov2_model_name(arch_name: str, patch_size: int) -> str:
  compact_arch_name = arch_name.replace("_", "")[:4]
  return f"dinov2_{compact_arch_name}{patch_size}"


def make_model(
  *,
  arch_name: str = "vit_large",
  img_size: int = 518,
  patch_size: int = 14,
  init_values: float = 1.0,
  ffn_layer: str = "mlp",
  block_chunks: int = 0,
  pretrained: bool = True,
  **kwargs,
):
  model_name = _make_dinov2_model_name(arch_name, patch_size)
  vit_kwargs = dict(
    img_size=img_size,
    patch_size=patch_size,
    init_values=init_values,
    ffn_layer=ffn_layer,
    block_chunks=block_chunks,
  )
  vit_kwargs.update(**kwargs)
  model = vits.__dict__[arch_name](**vit_kwargs)


  if pretrained:
    url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_pretrain.pth"
    state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
    model.load_state_dict(state_dict, strict=False)


  return model


model = make_model(arch_name="vit_small", pretrained=True)
model.eval()


dummy_input = torch.randn(1, 3, 518, 518, dtype=torch.float32)
torch.onnx.export(model, dummy_input, "dinov2_small_518.onnx", verbose=True,

使用onnxsim优化计算图，得到dinov2_small_518-sim.onnx模型

$ onnxsim dinov2_small_518.onnx dinov2_small_518-sim.onnx
Simplifying...
Finish! Here is the difference:
┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓
┃         ┃ Original Model ┃ Simplified Model ┃
┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩
│ Add       │ 135      │ 135       │
│ Concat     │ 2       │ 1        │
│ Constant    │ 143      │ 0        │
│ ConstantOfShape │ 1       │ 0        │
│ Conv      │ 1       │ 1        │
│ Div       │ 37       │ 37        │
│ Equal      │ 1       │ 0        │
│ Erf       │ 12       │ 12        │
│ Expand     │ 1       │ 0        │
│ Gather     │ 36       │ 36        │
│ MatMul     │ 72       │ 72        │
│ Mul       │ 86       │ 85        │
│ Pow       │ 25       │ 25        │
│ ReduceMean   │ 50       │ 50        │
│ Reshape     │ 25       │ 25        │
│ Shape      │ 1       │ 0        │
│ Slice      │ 2       │ 1        │
│ Softmax     │ 12       │ 12        │
│ Sqrt      │ 25       │ 25        │
│ Sub       │ 25       │ 25        │
│ Transpose    │ 37       │ 37        │
│ Where      │ 1       │ 0        │
│ Model Size   │ 85.6MiB    │ 85.4MiB     │
└─────────────────┴────────────────┴──────────────────┘

从onnxsim的输出log可见，采用ONNX模型格式支持DINOv2模型只需要支持Add、Concat、Conv、Div、Erf、Gather、MatMul、Mul、Pow、ReduceMean、Reshape、Slice、Softmax、Sqrt、Sub、Transpose算子即可，而这些算子均在AX650N的算子支持列表中。

4.2 模型编译

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

$ pulsar2 build --input model/dinov2_small_518-sim.onnx --output_dir dinov2 --config config/dinov2_config.json
Building onnx ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
patool: Extracting ./dataset/imagenet-32-images.tar ...
Transformer optimize level: 1
32 File(s) Loaded.
[10:57:11] AX Quantization Fusion Pass Running ...    Finished.
......
[11:05:52] AX Refine Int Parameter Pass Running ...    Finished.
Network Quantization Finished.
quant.axmodel export success: dinov2/quant/quant_axmodel.onnx
Building native ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
Building native ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
2023-06-15 11:06:21.367 | INFO   | yamain.command.build842 - unsupported ops: []
tiling op...  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 572/572 0:00:00
build op...  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1152/1152 0:00:13
......
2023-06-15 11:06:45.845 | INFO   | yasched.test_onepass1947 - max_cycle = 22347641
2023-06-15 11:06:57.014 | INFO   | yamain.command.build971 - QuantAxModel macs: 46696739328
2023-06-15 11:07:42.507 | INFO   | yamain.command.build:com

其中，config.json如下所示。本次示例单独展示config.json的目的是为了说明AX650N的工具链针对Transformer模型有特制的量化精度调优配置，开启后生成的模型精度更佳。

{
 "model_type": "ONNX",
 "npu_mode": "NPU3",
 "quant": {
  "input_configs": [
   {
    "tensor_name": "DEFAULT",
    "calibration_dataset": "./dataset/imagenet-32-images.tar",
    "calibration_size": 32,
    "calibration_mean": [123.675, 116.28, 103.53],
    "calibration_std": [58.395, 57.12, 57.375]
   }
  ],
  "calibration_method": "MSE",
  "transformer_opt_level": 1
 },
 "input_processors": [
  {
   "tensor_name": "DEFAULT",
   "tensor_format": "RGB",
   "src_format": "BGR",
   "src_dtype": "U8",
   "src_layout": "NHWC"
  }
 ],
 "compiler": {
  "check": 0
 }
}

五上板部署

5.1 AX-Samples

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

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

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

5.2 运行

由于DINOv2官方的Model ZOO只提供了基座模型，所以只能参考其PCA实现输出特征图可视化。

●“测试图片1”：DINOv2能够清晰的区分“小柴犬”的躯干、前爪、狗头、眼睛、鼻子、耳朵等语义信息。

●“测试图片2”：DINOv2能够清晰的区分图中小狗、自行车、护栏等语义信息。

5.3 性能统计

model	Input Size	AX650N推理耗时（ms）	帧率
ViT-S/14 distilled	518*518	28	35
ViT-B/14 distilled	518*518	92	10
ViT-L/14 distilled	518*518	305	3

六结束语

DINOv2作为一种无需Fine Tuning的自监督方法，在提取图像特征方面表现出色，适用于许多不同的视觉任务。期待更多基于DINOv2的下游任务出现。

审核编辑：汤梓红