OCR SOTA trên thiết bị với Core ML và dots.ocr

October 2, 2025
10 min read

SOTA OCR trên thiết bị với Core ML và dots.ocr

Mỗi năm phần cứng của chúng ta mạnh hơn một chút, các mô hình của chúng ta thông minh hơn một chút cho mỗi tham số. Vào năm 2025, việc chạy các mô hình thực sự cạnh tranh trên thiết bị trở nên khả thi hơn bao giờ hết. dots.ocr, một mô hình OCR 3B tham số từ RedNote, vượt qua Gemini 2.5 Pro trong OmniDocBench, biến OCR trở thành một trường hợp sử dụng trên thiết bị thực sự không thỏa hiệp. Việc chạy các mô hình trên thiết bị chắc chắn hấp dẫn đối với các nhà phát triển: không cần buôn lậu khóa API, không tốn chi phí và không cần mạng. Tuy nhiên, nếu chúng ta muốn các mô hình này chạy trên thiết bị, chúng ta cần lưu ý đến ngân sách tính toán và năng lượng hạn chế.

Đi vào Neural Engine, bộ tăng tốc AI tùy chỉnh của Apple đã được xuất xưởng với mọi thiết bị Apple kể từ năm 2017. Bộ tăng tốc này được thiết kế để có hiệu suất cao trong khi vẫn tiết kiệm pin. Một số thử nghiệm của chúng tôi đã cho thấy Neural Engine hiệu quả năng lượng hơn 12 lần so với CPU và hiệu quả năng lượng hơn 4 lần so với GPU.

Mặc dù tất cả những điều này nghe có vẻ rất hấp dẫn, nhưng thật không may, Neural Engine chỉ có thể truy cập được thông qua Core ML, khung ML mã nguồn đóng của Apple. Hơn nữa, ngay cả việc chuyển đổi một mô hình từ PyTorch sang Core ML cũng có thể gây ra một số thách thức và nếu không có mô hình được chuyển đổi trước hoặc một số kiến thức về các cạnh sắc, nó có thể gây khó khăn cho các nhà phát triển. May mắn thay, Apple cũng cung cấp MLX, một khung ML hiện đại và linh hoạt hơn nhắm mục tiêu đến GPU (không phải Neural Engine) và có thể được sử dụng kết hợp với Core ML.

Trong loạt bài gồm ba phần này, chúng tôi sẽ cung cấp một dấu vết lý luận về cách chúng tôi chuyển đổi dots.ocr để chạy trên thiết bị, sử dụng kết hợp CoreML và MLX. Quá trình này sẽ áp dụng được cho nhiều mô hình khác và chúng tôi hy vọng rằng điều này sẽ giúp làm nổi bật các ý tưởng và công cụ cần thiết cho các nhà phát triển đang tìm cách chạy các mô hình của riêng họ trên thiết bị.

Để theo dõi, hãy nhân bản repo. Bạn sẽ cần cài đặt uv và hf để chạy lệnh thiết lập:

bash ./boostrap.sh

Nếu bạn chỉ muốn bỏ qua và sử dụng mô hình đã chuyển đổi, bạn có thể tải xuống tại đây.

Chuyển đổi

Chuyển đổi từ PyTorch sang CoreML là một quá trình gồm hai bước:

Chụp biểu đồ thực thi PyTorch của bạn (thông qua torch.jit.trace hoặc phương pháp hiện đại hơn của torch.export).
Biên dịch biểu đồ đã chuyển đổi này thành .mlpackage bằng cách sử dụng coremltools.

Mặc dù chúng ta có một vài núm điều chỉnh mà chúng ta có thể tinh chỉnh cho bước 2, nhưng hầu hết quyền kiểm soát của chúng ta nằm ở bước 1, biểu đồ mà chúng ta cung cấp cho coremltools.

Tuân theo câu thần chú của lập trình viên về make it work, make it right, make it fast, trước tiên chúng ta sẽ tập trung vào việc làm cho quá trình chuyển đổi hoạt động trên GPU, ở FLOAT32 và với các hình dạng tĩnh. Khi chúng ta đã làm cho nó hoạt động, chúng ta có thể giảm độ chính xác và cố gắng chuyển sang Neural Engine.

Dots.OCR

Dots.OCR bao gồm hai thành phần chính: Một bộ mã hóa hình ảnh 1.2B tham số được đào tạo từ đầu, dựa trên kiến trúc NaViT và một xương sống Qwen2.5-1.5B. Chúng ta sẽ sử dụng CoreML để chạy bộ mã hóa hình ảnh và MLX để chạy xương sống LM.

Bước 0: Hiểu và đơn giản hóa mô hình

Để chuyển đổi một mô hình, tốt nhất là hiểu cấu trúc và chức năng trước khi bắt đầu. Nhìn vào tệp mô hình hóa hình ảnh gốc tại đây, chúng ta có thể thấy rằng bộ mã hóa hình ảnh tương tự như họ QwenVL. Giống như nhiều bộ mã hóa hình ảnh, bộ mã hóa hình ảnh cho dots hoạt động trên cơ sở vá, trong trường hợp này là các vá 14x14. Bộ mã hóa hình ảnh dots có khả năng xử lý video và hàng loạt hình ảnh. Điều này cho chúng ta cơ hội đơn giản hóa bằng cách chỉ xử lý một hình ảnh tại một thời điểm. Cách tiếp cận này thường xuyên được sử dụng trong các ứng dụng trên thiết bị, nơi chúng ta chuyển đổi một mô hình cung cấp các chức năng thiết yếu và lặp lại nếu chúng ta muốn xử lý nhiều hình ảnh.

Khi bắt đầu quá trình chuyển đổi, tốt nhất là bắt đầu với một mô hình khả thi tối thiểu. Điều này có nghĩa là loại bỏ bất kỳ chuông và còi nào không thực sự cần thiết để mô hình hoạt động. Trong trường hợp của chúng ta, dots có nhiều triển khai chú ý khác nhau có sẵn cho cả bộ mã hóa hình ảnh và xương sống LM. CoreML có rất nhiều cơ sở hạ tầng định hướng xung quanh toán tử scaled_dot_product_attention, mà họ đã giới thiệu trong iOS 18. Chúng ta có thể đơn giản hóa mô hình bằng cách loại bỏ tất cả các triển khai chú ý khác và chỉ tập trung vào sdpa đơn giản (không phải biến thể tiết kiệm bộ nhớ) cho bây giờ, cam kết tại đây.

Khi chúng ta đã thực hiện việc này, chúng ta sẽ thấy một thông báo cảnh báo đáng sợ khi chúng ta tải mô hình:

bash Sliding Window Attention is enabled but not implemented for sdpa; unexpected results may be encountered.

Mô hình không yêu cầu Sliding Window Attention để hoạt động, vì vậy chúng ta có thể vui vẻ tiếp tục.

Bước 1: Một bộ khung đơn giản

Sử dụng torch.jit.trace vẫn là phương pháp trưởng thành nhất để chuyển đổi mô hình sang CoreML. Chúng ta thường đóng gói nó trong một bộ khung đơn giản cho phép bạn sửa đổi các đơn vị tính toán được sử dụng và độ chính xác được chọn.

Bạn có thể xem bộ khung ban đầu tại đây. Nếu chúng ta chạy như sau trên triển khai mã gốc:

bash uv run convert.py –precision FLOAT32 –compute_units CPU_AND_GPU

Chúng ta sẽ gặp phải vấn đề đầu tiên (trong số rất nhiều).

Bước 2: Săn lỗi

Hiếm khi một mô hình sẽ chuyển đổi lần đầu tiên. Thông thường, bạn sẽ cần liên tục thực hiện các thay đổi xa hơn và xa hơn xuống biểu đồ thực thi cho đến khi bạn đến nút cuối cùng.

Vấn đề đầu tiên của chúng ta là lỗi sau:

bash ERROR - converting ‘outer’ op (located at: ‘vision_tower/rotary_pos_emb/192’): In op “matmul”, when x and y are both non-const, their dtype need to match, but got x as int32 and y as fp32

May mắn thay, lỗi này cung cấp cho chúng ta khá nhiều thông tin. Chúng ta có thể xem lớp VisionRotaryEmbedding và xem mã sau:

python def forward(self, seqlen: int) -> torch.Tensor: seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) freqs = torch.outer(seq, self.inv_freq) return freqs

Mặc dù torch.arange có một đối số dtype, coremltools bỏ qua điều này đối với arange và luôn xuất ra int32. Chúng ta có thể chỉ cần thêm một diễn viên sau arange để khắc phục vấn đề này, cam kết tại đây.

Sau khi khắc phục điều này, chạy lại quá trình chuyển đổi sẽ dẫn chúng ta đến vấn đề tiếp theo tại repeat_interleave:

bash ERROR - converting ‘repeat_interleave’ op (located at: ‘vision_tower/204’): Cannot add const [None]

Mặc dù lỗi này ít thông tin hơn, chúng ta chỉ có một lệnh gọi duy nhất đến repeat_interleave trong bộ mã hóa hình ảnh của chúng ta:

python cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( dim=0, dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, )

cu_seqlens được sử dụng để che giấu các chuỗi có độ dài biến đổi trong flash_attention_2. Nó được bắt nguồn từ tensor grid_thw, đại diện cho time, height và width. Vì chúng ta chỉ xử lý một hình ảnh duy nhất, chúng ta có thể chỉ cần loại bỏ lệnh gọi này, cam kết tại đây.

Vào lần tiếp theo! Lần này, chúng ta nhận được một lỗi khó hiểu hơn:

bash ERROR - converting ‘internal_op_tensor_inplace_fill’ op (located at: ‘vision_tower/0/attn/301_internal_tensor_assign_1’): _internal_op_tensor_inplace_fill does not support dynamic index

Điều này lại là do logic che giấu để xử lý các chuỗi có độ dài biến đổi. Vì chúng ta chỉ xử lý một hình ảnh duy nhất (không phải video hoặc hàng loạt hình ảnh), chúng ta không thực sự cần che giấu sự chú ý nào cả! Do đó, chúng ta có thể chỉ cần sử dụng một mặt nạ gồm tất cả True. Để chuẩn bị cho việc chuyển đổi Neural Engine, chúng ta cũng chuyển từ sử dụng mặt nạ boolean sang mặt nạ float gồm tất cả các số không, vì Neural Engine không hỗ trợ tensor bool cam kết tại đây

Với tất cả những điều này được thực hiện, bây giờ mô hình sẽ chuyển đổi thành công sang CoreML! Tuy nhiên, khi chúng ta chạy mô hình, chúng ta sẽ nhận được lỗi sau:

bash error: ‘mps.reshape’ op the result shape is not compatible with the input shape

Hình dạng lại này có thể ở nhiều nơi! May mắn thay, chúng ta có thể sử dụng một thông báo cảnh báo trước đó để giúp chúng ta theo dõi vấn đề:

bash TracerWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won’t change the number of iterations executed (and might lead to errors or silently give incorrect results). for t, h, w in grid_thw:

Hầu hết các trình biên dịch ML không thích luồng điều khiển động. May mắn cho chúng ta, vì chúng ta chỉ xử lý một hình ảnh duy nhất, chúng ta có thể chỉ cần loại bỏ vòng lặp và xử lý cặp h, w duy nhất, cam kết tại đây.

Và ở đó chúng ta có nó! Nếu chúng ta chạy lại quá trình chuyển đổi, chúng ta sẽ thấy rằng mô hình chuyển đổi thành công và phù hợp với độ chính xác PyTorch ban đầu:

bash Max difference: 0.006000518798828125, Mean difference: 1.100682402466191e-05

Bước 3: Điểm chuẩn

Bây giờ chúng ta đã có mô hình hoạt động, hãy đánh giá kích thước và hiệu suất. Tin tốt là mô hình đang hoạt động, tin xấu là nó hơn 5GB! Điều này hoàn toàn không thể chấp nhận được đối với việc triển khai trên thiết bị! Để đánh giá thời gian tính toán, chúng ta có thể sử dụng công cụ XCode tích hợp bằng cách gọi:

open DotsOCR_FLOAT32.mlpackage

điều này sẽ khởi chạy trình kiểm tra XCode cho mô hình. Sau khi nhấp vào + Performance Report và khởi chạy báo cáo trên tất cả các thiết bị tính toán, bạn sẽ thấy điều gì đó như sau:

Hơn một giây cho một lần chuyển tiếp duy nhất của bộ mã hóa hình ảnh! Chúng ta có rất nhiều việc phải làm.

Trong phần thứ hai của loạt bài này, chúng ta sẽ làm việc về sự tích hợp giữa CoreML và MLX, để chạy toàn bộ mô hình trên thiết bị. Trong phần thứ ba, chúng ta sẽ đi sâu vào các tối ưu hóa cần thiết để có được mô hình này chạy trên Neural Engine, bao gồm lượng tử hóa và hình dạng động.

AI Today - SkyAI