FormulaNet (Texo LaTeX OCR) - 昇腾 NPU 适配
模型介绍
FormulaNet 是基于 VisionEncoderDecoder 架构的 LaTeX OCR（光学公式识别）模型，是 Texo 项目的核心模型。Texo 是一个轻量级的开源 LaTeX OCR 模型，仅有 20M 参数，支持将公式图片转换为 LaTeX 代码。
- 模型名称: FormulaNet (Texo-distill)
- 原始模型地址:
- ModelScope: https://www.modelscope.cn/models/alephpi98/FormulaNet/files
- HuggingFace: https://huggingface.co/alephpi/FormulaNet
- GitHub: https://github.com/alephpi/Texo
- 任务类型: 公式识别 / LaTeX OCR（光学字符识别）
- 模型框架: PyTorch（VisionEncoderDecoder）
- 模型参数量: 20M（19,997,664）
- 输入格式: 公式图片（PNG/JPG，处理后为 3×384×384 张量）
- 输出格式: LaTeX 代码字符串
- 开源协议: AGPL-3.0
模型架构
FormulaNet 采用 VisionEncoderDecoder 架构：
- Encoder: HGNetv2（自定义 CNN 骨干网络，基于 PPFormulaNet）
- Decoder: mBART（轻量 Transformer 解码器，2层，hidden_size=384）
- Tokenizer: 自定义 687 词汇表的 PreTrainedTokenizerFast
该模型是 PPFormulaNet-S 的蒸馏版本，在 UniMERNet-1M 数据集上微调，保留了大部分性能的同时大幅减少了参数量。
依赖环境
系统要求
- Python >= 3.10
- Ascend NPU 驱动和固件（CANN 8.5.1+）
- PyTorch 2.0+
- torch_npu
Python 依赖
torch>=2.0.0
torchvision
transformers>=4.40.0
safetensors
Pillow>=10.0.0
numpy>=1.24.0NPU 适配说明
FormulaNet 基于标准 PyTorch + HuggingFace Transformers 框架，天然支持昇腾 NPU。适配过程包括：
- 自定义 Encoder 注册: 将 HGNetv2（model_type: my_hgnetv2）注册到 HuggingFace AutoModel 体系
- 图像预处理适配: 移除 OpenCV 依赖，使用纯 NumPy+Pillow 实现公式图像预处理（去边距、缩放填充、灰度归一化）
- 模型加载: 使用 safetensors 加载权重，适配 NPU 设备
环境准备
# 安装基础依赖
pip install torch torchvision transformers safetensors Pillow numpy
# 安装 torch_npu（根据环境选择合适的版本）
# pip install torch_npu推理命令
单张图像推理
python inference.py --image <input_image> --output result.txt --device cpu
python inference.py --image <input_image> --output result.txt --device npu批量推理
python inference.py --image_dir <image_dir> --output results.txt --device npu推理结果
测试图像
使用 Texo 项目提供的 5 张测试公式图像（包括单行公式、多行公式、复杂表格等）。
CPU 推理结果
| 图像 | 耗时(s) |
|---|---|
| 单行公式.png | 0.983 |
| 单行公式2.png | 1.885 |
| 多行公式.png | 5.249 |
| 复杂表格.png | 6.240 |
| 多行公式2.jpg | 2.194 |
| 总计 | 16.551 |
NPU 推理结果
| 图像 | 耗时(s) |
|---|---|
| 单行公式.png | 0.147 |
| 单行公式2.png | 0.411 |
| 多行公式.png | 1.272 |
| 复杂表格.png | 1.887 |
| 多行公式2.jpg | 0.430 |
| 总计 | 4.147 |
推理示例输出
输入图像: 单行公式.png
CPU 输出: 2^{2^{\lim_{n\to\infty}(\frac{1}{2})^n}}...
NPU 输出: 2^{2^{\lim_{n\to\infty}(\frac{1}{2})^n}}...
输入图像: 单行公式2.png
CPU 输出: \Re\left(\frac{-2i\kappa_0}{2\pi}\theta\right)...
NPU 输出: \Re\left(\frac{-2i\kappa_0}{2\pi}\theta\right)...CPU/NPU 精度测试方法
- 在 CPU 上运行推理，保存结果
- 在 NPU 上运行推理，保存结果
- 对比 CPU 和 NPU 的输出 LaTeX 代码（逐字符比较）
# CPU 推理
python inference.py --image_dir <test_images> --output cpu_results.txt --device cpu
# NPU 推理
python inference.py --image_dir <test_images> --output npu_results.txt --device npu
# 对比
python compare_cpu_npu.py --cpu cpu_results.txt --npu npu_results.txtCPU/NPU 精度测试结果
| 图像名称 | CPU耗时(s) | NPU耗时(s) | 加速比 | 完全一致 |
|---|---|---|---|---|
| 单行公式.png | 0.983 | 0.147 | 6.69x | 是 |
| 单行公式2.png | 1.885 | 0.411 | 4.59x | 是 |
| 多行公式.png | 5.249 | 1.272 | 4.13x | 是 |
| 复杂表格.png | 6.240 | 1.887 | 3.31x | 是 |
| 多行公式2.jpg | 2.194 | 0.430 | 5.10x | 是 |
精度测试结论
- 测试样本数: 5
- 完全一致: 5/5（100%）
- 平均字符级准确率: 100.0000%
- NPU 与 CPU 推理结果误差为 0.0000%，符合精度误差小于 1% 的要求
NPU 与 CPU 的推理结果完全一致（5/5 图像输出 LaTeX 代码完全相同），精度误差为 0%。
性能测试结果
| 指标 | CPU | NPU | 加速比 |
|---|---|---|---|
| 总耗时 (5张) | 16.551s | 4.147s | 3.99x |
| 平均耗时/张 | 3.310s | 0.829s | 3.99x |
| 最快单张 | 0.983s | 0.147s | 6.69x |
| 最慢单张 | 6.240s | 1.887s | 3.31x |
NPU（Ascend910）相比 CPU 实现了约 4 倍的平均推理加速。
部署和推理方法
核心推理代码
import torch
from PIL import Image, ImageOps
import numpy as np
from transformers import VisionEncoderDecoderModel, AutoTokenizer
from safetensors.torch import load_file
# 注册自定义 encoder
from texo.model.hgnet2 import HGNetv2Config, HGNetv2
from transformers import AutoConfig, AutoModel
AutoConfig.register("my_hgnetv2", HGNetv2Config)
AutoModel.register(HGNetv2Config, HGNetv2)
# 加载模型
model = VisionEncoderDecoderModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.to(device)
model.eval()
# 图像预处理
image = Image.open(image_path).convert('RGB')
image = crop_margin(image)
image = resize_with_padding(image, 384)
image = image.convert('L')
arr = np.array(image, dtype=np.float32) / 255.0
arr = (arr - 0.7931) / 0.1738
tensor = torch.from_numpy(arr).unsqueeze(0).float()
tensor = tensor.repeat(3, 1, 1).unsqueeze(0)
# 推理
with torch.no_grad():
outputs = model.generate(pixel_values=tensor.to(device), max_length=512)
pred = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
print(pred)完整推理脚本
参考 inference.py：
python inference.py --image formula.png --output result.txt --device npu推理成功证据
以下日志展示了 NPU 推理成功的关键信息：
NPU: Ascend910 (torch_npu 2.9.0.post1)模型标签
- #+NPU
- #+CV
- #+OCR
- #+公式识别
- #+LaTeX
- #+昇腾
- #+PyTorch
- #+生成式模型
推理截图
文件清单
inference.py- 推理脚本（支持 CPU/NPU）compare_cpu_npu.py- CPU/NPU 精度对比脚本requirements.txt- Python 依赖README.md- 本文档screenshot.png- 推理运行截图cpu_results.txt- CPU 推理结果npu_results.txt- NPU 推理结果comparison_results.txt- CPU/NPU 对比结果精度测试结果.log- 精度测试日志