Google ra mắt LiteRT.js: Chạy mô hình AI trực tiếp trên trình duyệt với sức mạnh WebGPU

Google đã ra mắt LiteRT.js, một thư viện JavaScript binding của LiteRT. LiteRT là thư viện suy luận (inference) trên thiết bị của Google, trước đây được gọi là TensorFlow Lite.

LiteRT.js chạy các mô hình.tflite trực tiếp bên trong trình duyệt. Vì quá trình suy luận diễn ra cục bộ, Google cho biết điều này giúp tăng cường quyền riêng tư của người dùng, không tốn chi phí máy chủ và có độ trễ cực thấp.

LiteRT.js là gì?

Đây không phải là một định dạng mô hình mới. Thay vào đó, Google đã biên dịch runtime gốc (native runtime) hiện có của họ sang WebAssembly và cung cấp cho JavaScript.

Các giải pháp AI trên web trước đây, bao gồm TensorFlow.js, dựa vào các nhân (kernel) viết bằng JavaScript. Google mô tả những giải pháp đó có hiệu năng thấp hơn. Ngược lại, LiteRT.js cung cấp runtime đa nền tảng gốc với các tối ưu hóa được giữ nguyên.

Nhờ đó, các ứng dụng web được thừa hưởng những thành quả đã đạt được ở các nền tảng khác. Các bản nâng cấp hiệu năng, cải tiến về lượng tử hóa (quantization) và tối ưu hóa phần cứng được xây dựng cho Android, iOS và máy tính để bàn cũng được áp dụng trên web.

Cách thức hoạt động: Một Runtime, ba Backend

Dưới runtime đó, LiteRT.js nhắm mục tiêu đến ba backend:

Hai quy tắc liên quan chi phối việc điều phối (dispatch). Thứ nhất, LiteRT.js không hỗ trợ ủy quyền một phần (partial delegation). Một đồ thị (graph) không thể bị chia tách giữa CPU và GPU.

Thứ hai, việc ủy quyền (delegation) là "được ăn cả, ngã về không" cho mỗi mô hình. Nếu một mô hình không thể được ủy quyền hoàn toàn cho bộ tăng tốc (accelerator) đã chọn, LiteRT sẽ quay lại thực thi bằng wasm. Đường dẫn CPU có độ bao phủ toán tử rộng nhất.

Hiệu năng

Với các backend đó, nhóm Google báo cáo hai kết quả khác biệt.

So với các runtime web khác, LiteRT.js nhanh hơn tới 3 lần trong quá trình suy luận trên CPU và GPU. Con số này áp dụng cho các mô hình thị giác máy tính và xử lý âm thanh cổ điển.

So với việc tự thực thi trên CPU, GPU hoặc NPU mang lại tốc độ nhanh hơn từ 5–60 lần. Điều này áp dụng cho các tác vụ thời gian thực đòi hỏi khắt khe như theo dõi đối tượng và chuyển đổi âm thanh thành văn bản.

Cả hai bài kiểm tra hiệu năng đều được thực hiện trong môi trường trình duyệt được kiểm soát trên MacBook Pro 2024 với chip Apple Silicon M4. Google lưu ý rằng kết quả sẽ thay đổi tùy thuộc vào GPU cục bộ, hiện tượng giảm xung nhịp do nhiệt (thermal throttling) và tối ưu hóa trình điều khiển. Con số "10x" đang lan truyền cùng thời điểm ra mắt không xuất hiện trong thông báo chính thức.

Cách đưa mô hình PyTorch vào

LiteRT Torch chuyển đổi các mô hình PyTorch sang.tflite chỉ trong một bước.

Tuy nhiên, các điều kiện tiên quyết rất nghiêm ngặt. Mô hình của bạn phải có khả năng xuất bằng torch.export.export, nghĩa là phải tương thích với TorchDynamo-exportable. Nó không được chứa các nhánh điều kiện Python phụ thuộc vào giá trị tensor khi chạy (runtime). Nó cũng không được có kích thước đầu vào hoặc đầu ra động, bao gồm cả kích thước batch.

Về kích thước, AI Edge Quantizer sẽ cấu hình các lược đồ lượng tử hóa trên các lớp mô hình khác nhau. Các mô hình.tflite đã được huấn luyện sẵn cũng có sẵn trên Kaggle và cộng đồng LiteRT trên Hugging Face.

Quy trình tối giản

Sau khi chuyển đổi, mã runtime rất ngắn gọn. Đây là đường dẫn WebGPU, đã được xác minh với @litertjs/core v2.5.2:

Khối cuối cùng đó cần được lưu ý. LiteRT.js không tự động thu gom rác (garbage-collect) các tensor. Mỗi Tensor phải được xóa một cách rõ ràng, nếu không ứng dụng sẽ bị rò rỉ bộ nhớ thiết bị. Đoạn mã trong bài đăng thông báo của Google đã bỏ qua bước này.

WebNN cần thêm một cờ (flag) bổ sung. LiteRT.js yêu cầu JSPI, giúp kết nối việc lập lịch nhân đồng bộ với việc thăm dò thiết bị bất đồng bộ:

Trước khi viết phần tiền xử lý, hãy kiểm tra với các đầu vào giả (fake inputs). Chạy npm i @litertjs/model-tester, sau đó là npx model-tester. Nó sẽ chạy mô hình của bạn trên WebNN, WebGPU và CPU bằng cách sử dụng các đầu vào ngẫu nhiên. Sử dụng model.getInputDetails để đọc tên và hình dạng đầu vào.

Các trường hợp sử dụng kèm ví dụ

Các API đó hỗ trợ bốn bản demo mà Google đã phát hành khi ra mắt:

LiteRT.js so với TensorFlow.js

Những bản demo đó đặt ra một câu hỏi hiển nhiên cho các đội ngũ phát triển ML trên web hiện nay.

Quan trọng là, hai thư viện này không loại trừ lẫn nhau. Google định vị LiteRT.js như một sự thay thế cụ thể cho các Graph Models của TF.js, chứ không phải thay thế toàn bộ thư viện.

TensorFlow.js vẫn là công cụ được khuyến nghị cho các tác vụ tiền xử lý và hậu xử lý. Gói @litertjs/tfjs-interop truyền các tensor giữa chúng thông qua runWithTfjsTensors. Hãy tránh sử dụng tensor.dataSync, vì nó gây ra sự sụt giảm hiệu năng đáng kể trên backend WebGPU.

Giải thích tương tác

Nội dung nhúng bên dưới mô phỏng sáu giai đoạn của quy trình trên mỗi backend.

Những điểm chính cần lưu ý

Nguồn

Michal Sutter là một chuyên gia khoa học dữ liệu với bằng Thạc sĩ Khoa học Dữ liệu từ Đại học Padova. Với nền tảng vững chắc về phân tích thống kê, học máy và kỹ thuật dữ liệu, Michal xuất sắc trong việc chuyển đổi các tập dữ liệu phức tạp thành những thông tin chi tiết có thể hành động được.