Memory footprint: оценка и снижение потребления памяти

ONNX Runtime Web (ORT Web) предназначен для выполнения моделей машинного обучения в браузере или в Node.js с использованием JavaScript. Одним из ключевых аспектов эффективного применения ORT Web является управление потреблением памяти, особенно при работе с крупными моделями или на устройствах с ограниченными ресурсами. Понимание структуры потребления памяти позволяет оптимизировать выполнение модели и избежать проблем с производительностью или отказом выполнения.

В ORT Web память потребляется несколькими основными компонентами:

  1. Граф модели – структура слоев и узлов, которая хранится в виде внутреннего представления модели.
  2. Тензоры – основные единицы хранения данных. Каждый вход, выход и промежуточный результат модели представлен тензором, выделяющим память под данные конкретного типа и размера.
  3. Буферы для выполнения операций – временные массивы, используемые для вычислений. Размер и количество буферов зависят от оптимизаций, включенных в рантайм, и порядка выполнения узлов модели.

Оценка потребления памяти

Оценка объема занимаемой памяти включает несколько подходов:

  • Анализ структуры модели: определение типов данных тензоров и их размеров. Для каждого тензора память вычисляется по формуле:

[ = ]

Например, тензор float32[1, 224, 224, 3] требует:

[ 1 = 602,112 ]

  • Мониторинг использования памяти во время выполнения: браузеры предоставляют API performance.memory (Chrome, Edge), позволяющие отслеживать общую и доступную память. Для Node.js можно использовать process.memoryUsage().

  • Профилирование тензоров ORT Web: через объект OrtSession доступна информация о созданных тензорах и их размере. Некоторые обертки предоставляют методы для подсчета активных тензоров и их потребления памяти в реальном времени.


Снижение потребления памяти

Эффективное управление памятью требует комбинации стратегий:

Использование подходящих типов данных

  • float16 или int8 вместо float32: уменьшает размер тензоров в 2–4 раза. В ORT Web поддерживаются операции с float16 и int8 при соответствующем преобразовании модели с помощью инструментов оптимизации ONNX.
  • Выбор оптимального типа для входных данных: если точность float32 не критична, преобразование входов к float16 снижает расход памяти без значительного влияния на качество предсказаний.

Очистка неиспользуемых тензоров

  • Тензоры в ORT Web автоматически освобождаются сборщиком мусора JavaScript, но явное уничтожение объектов OrtTensor с помощью tensor.dispose() позволяет снизить нагрузку на память во время интенсивных вычислений.
  • Для промежуточных результатов, которые больше не нужны после выполнения слоя или узла, рекомендуется освобождать ссылки сразу.

Минимизация числа параллельных сессий

  • Каждая активная сессия OrtSession держит собственные копии весов и буферов. Одновременное использование нескольких моделей увеличивает нагрузку на память.
  • При необходимости работы с несколькими моделями стоит разгружать сессии после использования:
session.dispose();

Инкрементальная загрузка модели

  • Для больших моделей ORT Web поддерживает lazy loading отдельных слоев или частей графа, что позволяет не загружать весь граф сразу.
  • Этот подход снижает пиковое потребление памяти, особенно на мобильных устройствах.

Оптимизация графа модели

  • Использование инструментов ONNX для свёртки слоев, удаления неиспользуемых узлов и квантования позволяет уменьшить количество промежуточных тензоров.
  • Оптимизированная модель требует меньше временных буферов при выполнении, что напрямую снижает memory footprint.

Практические рекомендации

  1. Всегда профилировать память для целевой платформы перед развертыванием.
  2. Предпочитать тензоры с меньшей точностью, если точность предсказаний допускает.
  3. Явно освобождать ресурсы с помощью dispose() для тензоров и сессий.
  4. Разделять большие задачи на последовательные шаги, чтобы избежать резкого роста пикового потребления памяти.
  5. Использовать инструменты ONNX для квантования и оптимизации модели перед загрузкой в ORT Web.

Эти методы позволяют контролировать и снижать memory footprint в ORT Web, обеспечивая стабильное и эффективное выполнение моделей на различных устройствах.