使用 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 主要包含兩個步驟：使用 torch.jit.trace 或 torch.export 捕獲 PyTorch 執行圖，然後使用 coremltools 將該圖編譯為 .mlpackage。

簡化模型以供設備上使用

為確保轉換成功，模型透過移除非單圖像處理所必需的功能進行簡化：

• 單圖像處理：模型被修改為一次處理一張圖像，而非影片或批次，這是設備上應用的常見最佳化。
• 注意力實作：所有注意力變體皆被移除，改為使用 Core ML 在 iOS 18 中支援的標準 scaled_dot_product_attention（sdpa）運算子。
• 移除滑動視窗注意力：由於模型運作不需要 Sliding Window Attention，故將其停用，以避免與 sdpa 的實作衝突。

轉換過程中的技術錯誤修正

在視覺編碼器的轉換過程中遇到並解決了多項技術障礙：

• 資料型別不匹配：coremltools 會忽略 torch.arange 中的 dtype 參數，預設為 int32。透過加入顯式的型別轉換，以確保在矩陣乘法時與 fp32 張量相容，解決了此問題。
• repeat_interleave 問題：用於在 flash_attention_2 中遮蔽可變長度序列的 repeat_interleave 呼叫被移除，因為在單圖像處理時並不需要此操作。
• 遮蔽邏輯：為支援 Neural Engine，將布林遮蔽改為全零的浮點遮蔽，因為 Neural Engine 不支援 bool 張量。
• 動態控制流程：移除遍歷 grid_thw 的迴圈，以消除動態控制流程，該流程通常不被 ML 編譯器支援。

初步基準測試與效能

初始轉換結果顯示，模型與原始 PyTorch 的精度相符，最大差異為 0.006000518798828125，平均差異為 1.100682402466191e-05。然而，最初的 FLOAT32 版本模型大小超過 5GB，且在 GPU 上執行視覺編碼器的單次前向傳播需超過一秒，顯示需要進一步的量化與最佳化，以供正式的設備上部署。