Источники готовых ONNX-моделей: ONNX Model Zoo, Hugging Face, Timm

ONNX Runtime Web (ORT Web) предоставляет возможность интегрировать модели машинного обучения в веб-приложения, обеспечивая высокую производительность и кросс-платформенность. Ключевой задачей при работе с ORT Web является выбор и подготовка модели. Для этого используются специализированные репозитории готовых ONNX-моделей. Основные из них — ONNX Model Zoo, Hugging Face и Timm. Рассмотрим их особенности, структуру и возможности интеграции.


ONNX Model Zoo

ONNX Model Zoo — это официальный каталог моделей, совместимых с форматом ONNX. Он включает широкий спектр моделей для классификации изображений, обработки текста, сегментации и других задач машинного обучения.

Особенности:

  • Формат ONNX: каждая модель доступна в формате .onnx, что обеспечивает мгновенную совместимость с ORT Web.
  • Метаданные и документация: каждая модель снабжена описанием архитектуры, входов и выходов, а также примером использования.
  • Примеры и скрипты конвертации: для большинства моделей предоставлены скрипты для конвертации из PyTorch или TensorFlow в ONNX.

Применение с ORT Web:

import * as ort from 'onnxruntime-web';

async function loadModel() {
  const session = await ort.InferenceSession.create('resnet50.onnx');
  const tensor = new ort.Tensor('float32', inputData, [1, 3, 224, 224]);
  const output = await session.run({ input: tensor });
  console.log(output);
}

Модели из Model Zoo обычно оптимизированы для CPU и GPU, но для браузера важна поддержка WebAssembly (WASM) и WebGL, что ORT Web реализует по умолчанию.


Hugging Face

Hugging Face — популярная платформа для NLP и мультимодальных моделей. Многие модели на Hugging Face поддерживают экспорт в формат ONNX.

Особенности:

  • ONNX-модели для NLP и CV: библиотека transformers позволяет конвертировать модели в ONNX, а на сайте есть готовые файлы .onnx.
  • Разнообразие архитектур: BERT, GPT, ViT, DALL-E и другие модели доступны в ONNX-версии.
  • Оптимизация и ускорение: Hugging Face предоставляет рекомендации по оптимизации моделей под веб-окружение, включая снижение размера модели и использование динамического батчинга.

Интеграция с ORT Web:

  1. Скачать модель в формате ONNX из репозитория Hugging Face.
  2. Разместить модель на сервере или CDN.
  3. Использовать InferenceSession.create() для загрузки модели в браузер.
const session = await ort.InferenceSession.create('bert.onnx');
const inputIds = new ort.Tensor('int64', [101, 2054, 2003, 1996, 102], [1, 5]);
const result = await session.run({ input_ids: inputIds });
console.log(result);

Hugging Face обеспечивает возможность работы с большими моделями через chunked loading и lazy initialization, что критично для браузерных приложений.


Timm

Timm — библиотека PyTorch для моделей компьютерного зрения. Несмотря на то, что Timm напрямую не хранит ONNX-модели, она позволяет легко экспортировать модели в формат ONNX для последующего использования в ORT Web.

Особенности:

  • Широкий выбор архитектур: EfficientNet, ResNet, Vision Transformer и другие.
  • Поддержка предобученных весов: модели можно взять с готовыми весами и сразу конвертировать.
  • Гибкость и кастомизация: позволяет менять слои и размер входного тензора перед экспортом в ONNX.

Пример конвертации в ONNX:

import torch
import timm

model = timm.create_model('resnet50', pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "resnet50.onnx", opset_version=12, input_names=['input'], output_names=['output'])

После экспорта модель можно загрузить в ORT Web, используя стандартные методы:

const session = await ort.InferenceSession.create('resnet50.onnx');

Timm особенно удобен для экспериментов с кастомными архитектурами и оптимизированными под веб-мобильные модели, где важны размер и скорость инференса.


Ключевые аспекты работы с ONNX-моделями в вебе

  • Выбор модели: ориентироваться на задачу (классификация, NLP, CV) и размер модели для веб-окружения.
  • Оптимизация модели: использовать инструменты вроде onnxruntime-tools для сжатия и оптимизации.
  • Поддержка веб-технологий: ORT Web поддерживает WebAssembly и WebGL, что позволяет использовать CPU и GPU браузера.
  • Форматы входных и выходных данных: строго соблюдать типы и размерность тензоров, указанные в документации модели.
  • Размещение модели: хранение на CDN или сервере с быстрым доступом критично для загрузки в браузере.

Использование ONNX Model Zoo, Hugging Face и Timm обеспечивает широкий выбор готовых моделей, упрощает интеграцию в веб-приложения и позволяет гибко настраивать производительность и размер моделей под конкретные задачи.