使用 Core ML 和 dots.ocr 的最先进 OCR

Hugging Face 详细介绍了将 dots.ocr（RedNote 的 3B 参数 OCR 模型）转换为在设备上使用 Apple 的 Core ML 和 MLX 框架执行的过程。该模型之所以重要，是因为它在 OmniDocBench 基准测试中超越了 Gemini 2.5 Pro，实现了高性能 OCR，无需 API 密钥、网络连接或持续费用。

通过 Neural Engine 的硬件加速

Apple 的 Neural Engine（NE）为 AI 工作负载提供了比 CPU 和 GPU 更高能效的替代方案。测试表明，Neural Engine 的能效 比 CPU 高 12 倍，比 GPU 高 4 倍，使其成为功耗受限的设备端模型的理想目标。然而，Neural Engine 只能通过 Core ML 访问，这是一个闭源框架，通常需要从 PyTorch 进行复杂的转换过程。

模型架构：dots.ocr

dots.ocr 使用由两个主要组件组成的混合架构：

  • Vision Encoder：基于 NaViT 架构的 1.2B 参数编码器，采用从头训练。该组件通过 Core ML 运行。
  • LM Backbone：Qwen2.5-1.5B 主干，通过 MLX 框架运行。

转换过程：从 PyTorch 到 Core ML

将模型从 PyTorch 转换为 Core ML 包括两个主要步骤：捕获 PyTorch 执行图（使用 torch.jit.trace 或 torch.export），并使用 coremltools 将该图编译为 .mlpackage。

为设备端使用简化模型

为确保转换成功，模型通过移除对单图像处理非必要的特性进行简化：

  • Single Image Processing：模型被修改为一次处理一张图像，而非视频或批量，这是设备端应用的常见优化。
  • Attention Implementation：所有注意力变体均被移除，改用标准的 scaled_dot_product_attention（sdpa）算子，该算子在 iOS 18 的 Core ML 中受支持。
  • Removal of Sliding Window Attention：由于模型运行不需要 Sliding Window Attention，已将其禁用，以避免与 sdpa 的实现冲突。

转换过程中的技术错误修复

在 vision encoder 的转换过程中遇到并解决了多个技术难题：

  • Dtype Mismatches：coremltools 忽略了 torch.arange 中的 dtype 参数，默认使用 int32。通过添加显式类型转换来确保在矩阵乘法中与 fp32 张量兼容，从而修复了此问题。
  • Repeat Interleave Issues：用于在 flash_attention_2 中对可变长度序列进行掩码的 repeat_interleave 调用被移除，因为在单图像处理时并不需要。
  • Masking Logic：为支持 Neural Engine，布尔掩码被替换为全零的浮点掩码，因为 Neural Engine 不支持 bool 张量。
  • Dynamic Control Flow：移除了遍历 grid_thw 的循环，以消除动态控制流，而动态控制流通常不被 ML 编译器支持。

初始基准测试与性能

初始转换结果显示，模型在最大差异 0.006000518798828125、平均差异 1.100682402466191e-05 的情况下，保持了原始 PyTorch 的精度。然而，模型的初始 FLOAT32 版本大小超过 5GB，并且在 GPU 上对 vision encoder 进行一次前向传播需要超过一秒，这表明需要进一步的量化和优化，以实现生产级的设备端部署。

Sources