LiteRT.js превосходит TensorFlow.js

LiteRT.js с использованием WebGPU выполняет инференс MobileNetV2 за 0,41 мс, обеспечивая 2439 FPS на десктопном GPU M2 Max — это более чем в 25 раз быстрее, чем TensorFlow.js на том же оборудовании. Разрыв увеличивается на более крупных моделях, выводя машинное обучение в вебе на уровень производительности, который раньше был доступен только нативным приложениям.

Почему этот бенчмарк важен

Разработчики использовали TensorFlow.js для инференса на устройстве, обычно через WebGL backend. WebGL отображает тензорные операции на 2D-графический конвейер браузера; он работает везде, но не был рассчитан на массивный параллелизм, необходимый моделям глубокого обучения. WebGPU, графический API следующего поколения, предоставляет прямой доступ к вычислительным блокам GPU. LiteRT.js — это первая библиотека, предоставляющая движок инференса на базе WebGPU для моделей TensorFlow-Lite, и Google сообщает о трехкратном ускорении. Независимые тесты показывают еще больший прирост, что ставит вопрос: станет ли WebGPU стандартом для веб-ML?

Как проводилось тестирование

Мы протестировали две модели классификации изображений — MobileNetV2 и EfficientNet-Lite4, обе из которых были конвертированы в TensorFlow-Lite. Тесты проводились на чипе Apple Silicon M2 Max. Для каждой модели мы измеряли задержку (latency) одного прямого прохода (forward pass) на протяжении множества итераций и фиксировали количество кадров в секунду (FPS). Мы запустили те же модели с TensorFlow.js на WebGL backend и зафиксировали время «холодного старта» каждой библиотеки.

Результаты: скорость и запуск

  • MobileNetV2

    • LiteRT.js (WebGPU): задержка 0,41 мс → 2439 FPS
    • TensorFlow.js (WebGL): задержка 10,82 мс → 92 FPS

  • EfficientNet-Lite4

    • LiteRT.js (WebGPU): задержка 0,62 мс → 1623 FPS
    • TensorFlow.js (WebGL): данные не представлены, но MobileNetV2 уже демонстрирует преимущество более чем в 25 раз.

TensorFlow.js потребовалось около 10 секунд (10 014 мс) на компиляцию WebGL-шейдеров перед первым инференсом. LiteRT.js был готов за 4,3 мс, что практически устраняет задержку «холодного старта» для интерактивных приложений.

При увеличении размера модели в пять раз задержка выросла всего на 50%, что подтверждает: параллелизм GPU поглощает большую часть дополнительной нагрузки. В результате достигается такой уровень производительности, который позволяет браузеру обрабатывать видеопотоки в реальном времени, AR-слои или быстрое прототипирование моделей компьютерного зрения без обращения к серверу.

Что скрывается за цифрами

  • Batching (пакетирование) — большинство моделей TensorFlow-Lite фиксируют размер пакета (batch dimension) на значении 1. Подача нескольких изображений за один вызов не поддерживается «из коробки», что вынуждает разработчиков создавать Web Workers или вручную выстраивать конвейер входных данных.
  • Управление памятью — LiteRT.js не выполняет автоматическую сборку мусора для тензоров GPU. Разработчики должны вызывать .delete() для каждого созданного тензора, иначе после нескольких сотен инференсов можно столкнуться с нехваткой памяти GPU.
  • Доступность на оборудовании — WebGPU стабильно работает в десктопных Chrome и Firefox. Мобильные браузеры (включая Chrome на Android и Safari на iOS) предоставляют доступ к API либо через экспериментальные флаги, либо не предоставляют вовсе. На этих платформах единственным универсальным вариантом остается WASM backend, но он заметно медленнее работает с крупными моделями.

Эти ограничения означают, что, хотя чистая скорость впечатляет, инженерные усилия для её достижения могут быть весьма значительными.

Ограничения и компромиссы

WebGL backend в TensorFlow.js по-прежнему обеспечивает почти универсальную совместимость. Разработчик, ориентированный на смешанную аудиторию (десктоп, Android, iOS), может полагаться на единый путь выполнения кода, который работает везде, пусть и с меньшей пропускной способностью. WASM backend работает практически на любом оборудовании, но уступает WebGPU на крупных моделях, представленных в этом тесте.

LiteRT.js показывает себя наилучшим образом, когда целевой средой является десктоп с включенным WebGPU. Он запускает модели .tflite напрямую, позволяя разработчикам загружать модели с Hugging Face или Kaggle без конвертации, сохраняя оригинальное квантование и производительность. Компромиссом здесь является более узкая область применения и необходимость тщательного управления ресурсами.

На что стоит обратить внимание разработчикам

  • Целевая платформа — для десктопного веб-инструмента (например, приложения для дизайна с переносом стиля в реальном времени) WebGPU с LiteRT.js, скорее всего, будет лучшим выбором.
  • Размер модели — крупные, ресурсоемкие модели получают наибольшую выгоду от параллелизма GPU; использование малых моделей может не оправдать дополнительных инженерных затрат.
  • Дисциплина управления памятью — планируйте явное удаление тензоров или оборачивайте вызовы инференса в область видимости, которая автоматически освобождает ресурсы.
  • Стратегия отката (fallback) — предусмотрите использование WASM или WebGL для браузеров, которые не поддерживают WebGPU, чтобы приложение оставалось функциональным для более широкой аудитории.

Итог

LiteRT.js доказывает, что WebGPU позволяет перевести веб-инференс в субмиллисекундный диапазон, обеспечивая скорость более чем в 25 × выше, чем у TensorFlow.js на настольных GPU. Технология все еще находится на стадии становления, и разработчикам приходится учитывать ограничения пакетной обработки, необходимость ручной очистки памяти и ограниченную поддержку мобильных устройств. Для решений, ориентированных прежде всего на десктоп и требующих производительности в реальном времени, новая библиотека предлагает перспективный путь развития; для обеспечения кроссплатформенного охвата по-прежнему необходимы более старые бэкенды WebGL и WASM.