基于AX650N部署視覺(jué)大模型DINOv2

一背景

最近一段時(shí)間，基于Transformer網(wǎng)絡(luò)結(jié)構(gòu)的視覺(jué)大模型呈現(xiàn)出爆發(fā)式增長(zhǎng)，繼Segment Anything（SAM）之后，Meta AI再次發(fā)布重量級(jí)開(kāi)源項(xiàng)目——DINOv2。DINOv2可以抽取到強(qiáng)大的圖像特征，且在下游任務(wù)上不需要微調(diào)，這使得它適合作為許多不同的應(yīng)用中新的Backbone。

與之前發(fā)布的Segment Anything相比，DINOv2在應(yīng)用領(lǐng)域和適用范圍上更加廣泛，原論文中的實(shí)驗(yàn)也涵蓋了多個(gè)CV中經(jīng)典的下游任務(wù)。

二DINOv2簡(jiǎn)介

在Meta AI官方的Blog中，將DINOv2的特性總結(jié)如下：

● 自監(jiān)督訓(xùn)練

●無(wú)需Fine Tuning即可遷移至下游任務(wù)

●開(kāi)箱即用的視覺(jué)基座大模型

DINOv2是一種新的高精度計(jì)算機(jī)視覺(jué)模型訓(xùn)練方法，使用自監(jiān)督學(xué)習(xí)來(lái)實(shí)現(xiàn)與該領(lǐng)域中使用的標(biāo)準(zhǔn)方法相匹配。與其他自監(jiān)督系統(tǒng)一樣，使用DINOv2方法的模型可以在不需要任何相關(guān)原數(shù)據(jù)的情況下對(duì)任何圖像集合進(jìn)行訓(xùn)練。這意味著它可以從它所接收到的所有圖像中學(xué)習(xí)，而不僅僅是那些包含特定一組標(biāo)簽或標(biāo)題的圖像。DINOv2提供了可直接用作簡(jiǎn)單線性分類器輸入的高性能特征。這種靈活性意味著DINOv2可用于創(chuàng)建許多不同計(jì)算機(jī)視覺(jué)任務(wù)的多用途骨干。

Visualization of PCA

論文中的實(shí)驗(yàn)展示了DINOv2在下游任務(wù)上的出色能力，例如分類、分割和圖像檢索等應(yīng)用領(lǐng)域。其中，最令人驚訝的是，在深度估計(jì)方面，DINOv2的結(jié)果明顯優(yōu)于in-domain與out-of-domain的SOTA的pipeline。作者認(rèn)為這種強(qiáng)大的域外表現(xiàn)是自監(jiān)督特征學(xué)習(xí)和輕量級(jí)任務(wù)特定模塊（例如線性分類器）相結(jié)合的結(jié)果。

三AX650N

AX650N是一款兼具高算力與高能效比的SoC芯片，集成了八核Cortex-A55 CPU，10.8TOPs@INT8 NPU（針對(duì)Transformer模型進(jìn)行了定制優(yōu)化），支持8K@30fps的ISP，以及H.264、H.265編解碼的VPU。接口方面，AX650N支持64bit LPDDR4x，多路MIPI輸入，千兆Ethernet、USB、以及HDMI 2.0b輸出，并支持32路1080p@30fps解碼。強(qiáng)大的性能可以讓AX650N幫助用戶在智慧城市、智慧教育、智能制造等領(lǐng)域發(fā)揮更大的價(jià)值。

四模型轉(zhuǎn)換

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

4.1 模型下載

●下載DINOv2的開(kāi)源項(xiàng)目

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

●進(jìn)入dinov2目錄，使用下列文件替換 dinov2/models/vision_transformer.py

vision_transformer.py可從我們的開(kāi)源項(xiàng)目中獲取

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

●獲取ONNX模型，在dinov2根目錄下執(zhí)行

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優(yōu)化計(jì)算圖，得到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可見(jiàn)，采用ONNX模型格式支持DINOv2模型只需要支持Add、Concat、Conv、Div、Erf、Gather、MatMul、Mul、Pow、ReduceMean、Reshape、Slice、Softmax、Sqrt、Sub、Transpose算子即可，而這些算子均在AX650N的算子支持列表中。

4.2 模型編譯

使用AX650N配套的AI工具鏈Pulsar2，一鍵完成圖優(yōu)化、離線量化、編譯、對(duì)分功能。

$ 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如下所示。本次示例單獨(dú)展示config.json的目的是為了說(shuō)明AX650N的工具鏈針對(duì)Transformer模型有特制的量化精度調(diào)優(yōu)配置，開(kāi)啟后生成的模型精度更佳。

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

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

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

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

5.2 運(yùn)行

由于DINOv2官方的Model ZOO只提供了基座模型，所以只能參考其PCA實(shí)現(xiàn)輸出特征圖可視化。

●“測(cè)試圖片1”：DINOv2能夠清晰的區(qū)分“小柴犬”的軀干、前爪、狗頭、眼睛、鼻子、耳朵等語(yǔ)義信息。

●“測(cè)試圖片2”：DINOv2能夠清晰的區(qū)分圖中小狗、自行車、護(hù)欄等語(yǔ)義信息。

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

六結(jié)束語(yǔ)

DINOv2作為一種無(wú)需Fine Tuning的自監(jiān)督方法，在提取圖像特征方面表現(xiàn)出色，適用于許多不同的視覺(jué)任務(wù)。期待更多基于DINOv2的下游任務(wù)出現(xiàn)。