LiteRT.js đang vượt mặt TensorFlow.js

LiteRT.js với WebGPU thực hiện suy luận MobileNetV2 trong 0,41 ms, đạt tốc độ 2.439 FPS trên GPU máy tính để bàn M2 Max—nhanh hơn hơn 25 lần so với TensorFlow.js trên cùng một phần cứng. Khoảng cách này còn nới rộng hơn đối với các mô hình lớn hơn, đưa học máy trên nền tảng web lên một phân khúc hiệu năng vốn trước đây chỉ dành cho các ứng dụng native.

Tại sao kết quả benchmark này lại quan trọng

Các nhà phát triển thường sử dụng TensorFlow.js để suy luận trên thiết bị, thường là thông qua backend WebGL. WebGL ánh xạ các phép toán tensor lên pipeline đồ họa 2D của trình duyệt; nó hoạt động ở mọi nơi nhưng không được thiết kế cho khả năng tính toán song song khổng lồ mà các mô hình học sâu cần. WebGPU, một API đồ họa thế hệ tiếp theo, cho phép truy cập trực tiếp vào các đơn vị tính toán của GPU. LiteRT.js là thư viện đầu tiên cung cấp công cụ suy luận hỗ trợ WebGPU cho các mô hình TensorFlow-Lite, và Google báo cáo tốc độ tăng gấp ba lần. Các thử nghiệm độc lập cho thấy mức tăng còn lớn hơn nữa, đặt ra câu hỏi liệu WebGPU có trở thành con đường mặc định cho web ML hay không.

Cách thức thực hiện thử nghiệm

Chúng tôi đã thực hiện benchmark hai mô hình phân loại hình ảnh—MobileNetV2 và EfficientNet-Lite4—cả hai đều được chuyển đổi sang TensorFlow-Lite. Các thử nghiệm được chạy trên chip Apple-silicon M2 Max. Đối với mỗi mô hình, chúng tôi đo độ trễ của một lượt truyền thẳng (forward pass) duy nhất qua nhiều lần lặp và báo cáo số khung hình trên giây (FPS). Chúng tôi cũng chạy các mô hình tương tự với TensorFlow.js trên backend WebGL và ghi lại thời gian khởi động lạnh (cold-start time) của mỗi thư viện.

Kết quả: tốc độ và khởi động

  • MobileNetV2

    • LiteRT.js (WebGPU): 0,41 ms độ trễ → 2.439 FPS
    • TensorFlow.js (WebGL): 10,82 ms độ trễ → 92 FPS

  • EfficientNet-Lite4

    • LiteRT.js (WebGPU): 0,62 ms độ trễ → 1.623 FPS
    • TensorFlow.js (WebGL): không được báo cáo, nhưng MobileNetV2 đã cho thấy lợi thế hơn 25 lần.

TensorFlow.js cần khoảng 10 giây (10.014 ms) để biên dịch các WebGL shader trước khi thực hiện suy luận đầu tiên. LiteRT.js đã sẵn sàng chỉ trong 4,3 ms, về cơ bản loại bỏ được sự chậm trễ do khởi động lạnh đối với các ứng dụng tương tác.

Khi chúng tôi tăng kích thước mô hình lên gấp 5 lần, độ trễ chỉ tăng 50%, xác nhận rằng khả năng tính toán song song của GPU đã hấp thụ hầu hết khối lượng công việc tăng thêm. Kết quả là một ranh giới hiệu năng mới cho phép trình duyệt xử lý các luồng video thời gian thực, các lớp phủ AR, hoặc tạo nguyên mẫu nhanh các mô hình thị giác mà không cần gửi dữ liệu về máy chủ.

Những gì các con số chưa nói lên

  • Batching (Gom nhóm) – Hầu hết các mô hình TensorFlow-Lite đều cố định kích thước batch là 1. Việc đưa nhiều hình ảnh vào trong một lần gọi không được hỗ trợ sẵn, buộc các nhà phát triển phải khởi tạo Web Workers hoặc tự thiết lập pipeline cho đầu vào một cách thủ công.
  • Quản lý bộ nhớ – LiteRT.js không tự động thu gom rác (garbage-collect) cho các tensor trên GPU. Các nhà phát triển phải gọi .delete() trên mỗi tensor mà họ tạo ra, nếu không sẽ có nguy cơ làm cạn kiệt bộ nhớ GPU sau vài trăm lần suy luận.
  • Khả năng tiếp cận phần cứng – WebGPU hoạt động ổn định trên Chrome và Firefox phiên bản máy tính để bàn. Các trình duyệt di động—bao gồm Chrome trên Android và Safari trên iOS—chỉ cung cấp API này thông qua các cờ thử nghiệm (experimental flags) hoặc thậm chí không hỗ trợ. Trên các nền tảng đó, backend WASM vẫn là con đường duy nhất có sẵn phổ biến, nhưng nó chậm hơn đáng kể đối với các mô hình lớn hơn.

Những hạn chế này có nghĩa là mặc dù tốc độ thô rất ấn tượng, nhưng nỗ lực kỹ thuật để đạt được nó có thể không hề đơn giản.

Hạn chế và sự đánh đổi

Backend WebGL của TensorFlow.js vẫn cung cấp khả năng tương thích gần như toàn diện. Một nhà phát triển hướng tới đối tượng người dùng hỗn hợp—máy tính để bàn, Android, iOS—có thể dựa vào một luồng mã duy nhất chạy được ở mọi nơi, mặc dù với thông lượng thấp hơn. Backend WASM chạy trên hầu hết mọi phần cứng nhưng tụt hậu so với WebGPU đối với các mô hình lớn hơn được đo lường ở đây.

LiteRT.js tỏa sáng khi mục tiêu là môi trường máy tính để bàn đã bật WebGPU. Nó chạy trực tiếp các mô hình .tflite, cho phép các nhà phát triển lấy mô hình từ Hugging Face hoặc Kaggle mà không cần chuyển đổi, giúp bảo toàn nguyên vẹn quá trình lượng tử hóa (quantization) và hiệu suất ban đầu. Sự đánh đổi là phạm vi triển khai hẹp hơn và cần phải quản lý tài nguyên cẩn thận.

Những điều nhà phát triển nên cân nhắc

  • Nền tảng mục tiêu – Đối với một công cụ máy tính để bàn chỉ chạy trên web (ví dụ: một ứng dụng thiết kế áp dụng chuyển đổi phong cách - style transfer - trong thời gian thực), WebGPU với LiteRT.js có khả năng là lựa chọn tốt nhất.
  • Kích thước mô hình – Các mô hình lớn hơn, nặng về tính toán sẽ hưởng lợi nhiều nhất từ khả năng tính toán song song của GPU; các mô hình nhỏ hơn có thể không xứng đáng với nỗ lực kỹ thuật tăng thêm.
  • Kỷ luật bộ nhớ – Hãy lập kế hoạch xóa tensor một cách rõ ràng hoặc bao bọc các lệnh gọi suy luận trong một phạm vi (scope) có khả năng tự động giải phóng tài nguyên.
  • Chiến lược dự phòng (Fallback) – Cung cấp phương án dự phòng WASM hoặc WebGL cho các trình duyệt không thể bật WebGPU, nhằm giữ cho ứng dụng hoạt động được với đối tượng người dùng rộng hơn.

Lời kết

LiteRT.js chứng minh rằng WebGPU có thể đưa quá trình suy luận trên web vào ngưỡng dưới một mili giây, mang lại tốc độ nhanh hơn 25 × so với TensorFlow.js trên các GPU máy tính để bàn. Công nghệ này vẫn đang trong quá trình hoàn thiện, và các nhà phát triển phải đối mặt với các giới hạn về batching, việc dọn dẹp bộ nhớ thủ công và sự hỗ trợ hạn chế trên thiết bị di động. Đối với các trải nghiệm ưu tiên máy tính để bàn đòi hỏi hiệu suất thời gian thực, thư viện mới này mở ra một hướng đi đầy hứa hẹn; còn để tiếp cận đa nền tảng, các backend WebGL và WASM cũ hơn vẫn đóng vai trò thiết yếu.