Prefill сессий и переиспользование между вызовами

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

Создание сессии

Сессия в ONNX Runtime Web создаётся с использованием класса InferenceSession. При создании сессии модель загружается, компилируется под целевой движок (WebAssembly или WebGPU) и готовится к инференсу. Типичная инициализация выглядит так:

import * as ort from "onnxruntime-web";

const session = await ort.InferenceSession.create("model.onnx", {
  executionProviders: ["wasm"], // или "webgl", "webgpu"
  graphOptimizationLevel: "all"
});

Ключевые параметры:

  • executionProviders – определяет движок, на котором будет выполняться инференс.
  • graphOptimizationLevel – уровень оптимизации графа модели, напрямую влияющий на скорость выполнения.

Prefill сессии

Prefill сессии означает предварительное создание и загрузку модели до фактического вызова инференса. В браузере это критически важно для уменьшения задержек при первом запросе. Если модель большая, загрузка и компиляция могут занимать заметное время, поэтому предзагрузка сессии позволяет сразу обрабатывать запросы после инициализации.

let prefilledSession = null;

async function prefillSession() {
  if (!prefilledSession) {
    prefilledSession = await ort.InferenceSession.create("model.onnx", {
      executionProviders: ["webgpu"],
      graphOptimizationLevel: "all"
    });
  }
  return prefilledSession;
}

Преимущества Prefill:

  • Устраняет задержку первого инференса.
  • Позволяет асинхронно подготовить несколько моделей параллельно.
  • Снижает нагрузку на основной поток, если используется WebWorker.

Переиспользование сессий между вызовами

После того как сессия создана, её можно многократно использовать для разных запросов. Это особенно важно для веб-приложений, где модель может обрабатывать сотни или тысячи инференсов в течение жизни страницы.

async function runInference(inputData) {
  const session = await prefillSession();
  
  const feeds = { input: new ort.Tensor("float32", inputData, [1, 3, 224, 224]) };
  const results = await session.run(feeds);
  
  return results.output.data;
}

Особенности переиспользования:

  • Сессия остаётся в памяти, что экономит ресурсы при повторных вызовах.
  • Не требуется повторная компиляция модели, что критично для больших сетей.
  • Можно безопасно использовать одну сессию в разных асинхронных потоках, если не модифицируется состояние модели между вызовами.

Работа с несколькими моделями

В случае наличия нескольких моделей можно заранее создавать prefill сессий для каждой из них и хранить их в объекте или Map:

const sessions = new Map();

async function getSession(modelName) {
  if (!sessions.has(modelName)) {
    const session = await ort.InferenceSession.create(modelName, { executionProviders: ["wasm"] });
    sessions.set(modelName, session);
  }
  return sessions.get(modelName);
}

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

Оптимизация памяти и потоков

  • Использование dispose() для Tensor’ов освобождает память после инференса.
  • Для WebAssembly провайдеров полезно предварительно выделять буферы и переиспользовать их, чтобы снизить нагрузку на сборщик мусора.
  • WebGPU провайдер автоматически управляет памятью видеокарты, но переиспользование сессий снижает частоту выделения GPU-ресурсов.

Асинхронное управление сессиями

Для приложений с большим количеством одновременных запросов рекомендуется использовать очередь запросов к сессии. Это позволяет избежать гонок при инициализации и обеспечивает стабильное время отклика:

const requestQueue = [];

async function enqueueInference(inputData) {
  return new Promise((resolve) => {
    requestQueue.push({ inputData, resolve });
  });
}

async function processQueue() {
  while (true) {
    if (requestQueue.length > 0) {
      const { inputData, resolve } = requestQueue.shift();
      const result = await runInference(inputData);
      resolve(result);
    } else {
      await new Promise(r => setTimeout(r, 10));
    }
  }
}

processQueue();

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

Выводы по архитектуре

Использование prefill сессий и переиспользование между вызовами является одной из основных практик для производительных веб-приложений на базе ONNX Runtime Web. Это снижает задержки, экономит память и ресурсы процессора, упрощает работу с несколькими моделями и обеспечивает предсказуемое время отклика.

Благодаря этому подходу даже большие модели с тяжелыми вычислительными графами становятся практичными для использования в браузере или в Node.js окружении, обеспечивая масштабируемый и стабильный инференс.