WebSSL-MAE-700M-FULL2B-224 Ascend NPU 部署指南

项目简介

WebSSL-MAE-700M-FULL2B-224 是一个 632M 参数的 Vision Transformer (ViT-L) 模型，使用 Masked Autoencoder (MAE) 自监督学习在 20亿张网络图片上训练。本项目提供其在华为 Ascend NPU 环境下的部署方案。

特性

• 支持 Ascend NPU 推理加速
• CPU vs NPU 精度对比测试 (< 1% 误差)
• 224x224 分辨率输入
• 257 个 patch tokens (16x16 + 1 cls token)

环境要求

• 硬件: 华为 Ascend 910 系列 NPU
• CANN: 8.0.RC1 或更高版本
• PyTorch: 2.0+ with torch_npu
• Docker: 容器名称 test-modelagent

目录结构

/data/ysws/agentsp/webssl-mae700m-full2b-224-ascend/
├── inference.py          # 精度测试脚本
├── log.txt               # 测试日志
├── README.md             # 本文档
└── webssl_teaser.png     # 模型架构图

部署步骤

1. 进入容器

docker exec -it test-modelagent bash

2. 设置环境变量

source /usr/local/Ascend/ascend-toolkit/set_env.sh

3. 准备模型文件

模型文件应放在 /data/ysws/agentsp/webssl-mae700m-full2b-224/ 目录下：

• model.safetensors - 模型权重 (632M 参数, 2.5GB)
• config.json - 模型配置
• preprocessor_config.json - 预处理配置
• webssl_teaser.png - 模型架构图

4. 执行精度测试

cd /data/ysws/agentsp/webssl-mae700m-full2b-224-ascend/
python3 inference.py --precision_test

5. 运行推理

cd /data/ysws/agentsp/webssl-mae700m-full2b-224-ascend/
python3 inference.py

测试验证

精度测试结果

性能数据

测试日志

完整测试日志保存在 log.txt

使用示例

运行推理

import torch
from pathlib import Path
from safetensors.torch import load_file

model_path = "/data/ysws/agentsp/webssl-mae700m-full2b-224"

# 加载模型权重
state_dict = load_file(f"{model_path}/model.safetensors")

# 移动到 NPU
tensor = state_dict["cls_token"].to("npu:0")

使用 transformers

from transformers import AutoImageProcessor, ViTModel
from PIL import Image

processor = AutoImageProcessor.from_pretrained("facebook/webssl-mae700m-full2b-224")
model = ViTModel.from_pretrained("facebook/webssl-mae700m-full2b-224").cuda().eval()

image = Image.open("path/to/image.jpg")
inputs = processor(images=image, return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model(**inputs)

encoder_hidden_states = outputs.last_hidden_state

模型结构

WebSSL-MAE-700M-FULL2B-224 是一个 ViT-Large 模型 (比 300M 版本更大)：

每层 ViTBlock 包含：

• LayerNorm
• ViTAttention (16 heads, dim=1280)
• LayerNorm
• ViTMLP (intermediate=5120, GELU activation)

输入输出格式

• 输入: (B, 3, 224, 224) - RGB 图片
• 输出: (B, 257, 1280) - [CLS] token + 256 patch embeddings

常见问题

Q: 精度测试失败?

A: 检查 NPU 驱动是否正确安装，确保 CANN 环境变量已 source。

Q: 模型加载慢?

A: 632M 参数模型文件较大 (~2.5GB)，safetensors 解析需要一些时间。

Q: 推理时间较长?

A: 632M 参数完整推理约 6 秒属于正常范围，可使用批处理提升吞吐量。

许可证

本项目遵循 WebSSL-MAE-700M-FULL2B-224 原始许可证 (cc-by-nc-4.0)。