Профилирование с tf.profile

TensorFlow.js предоставляет мощные инструменты для анализа производительности моделей и операций в браузере или Node.js. Одним из ключевых механизмов является функция tf.profile, позволяющая детально отслеживать потребление памяти, количество операций и другие показатели во время выполнения кода.

Основные возможности tf.profile

tf.profile — это асинхронная функция, которая принимает колбэк и возвращает объект с подробной статистикой. Основные показатели, доступные через tf.profile:

  • newBytes — количество байт, выделенных за время выполнения профилируемой функции.
  • newTensors — количество созданных тензоров.
  • kernels — массив объектов, содержащих информацию о каждой операции (ядре), выполненной во время профилирования: тип операции, размер входных и выходных тензоров, затраченное время.
  • peakBytes — максимальное потребление памяти в процессе выполнения.
  • result — результат, возвращаемый профилируемой функцией.

Пример использования:

import * as tf from '@tensorflow/tfjs';

async function profileExample() {
  const profile = await tf.profile(() => {
    const a = tf.randomNormal([1000, 1000]);
    const b = tf.randomNormal([1000, 1000]);
    return a.matMul(b);
  });

  console.log('Количество новых тензоров:', profile.newTensors);
  console.log('Выделено памяти (байт):', profile.newBytes);
  console.log('Пиковое потребление памяти (байт):', profile.peakBytes);
  console.log('Типы операций:', profile.kernels.map(k => k.name));
  profile.result.dispose();
}

profileExample();

Работа с памятью

Одной из частых проблем при разработке на TensorFlow.js является утечка памяти. tf.profile позволяет выявлять, какие тензоры остаются в памяти после выполнения операций.

  • newTensors помогает определить, сколько объектов было создано.
  • peakBytes показывает пиковое потребление памяти, что важно для крупных моделей.
  • После анализа рекомендуется вручную освобождать тензоры через dispose().
const profile = await tf.profile(() => {
  const x = tf.ones([5000, 5000]);
  const y = tf.ones([5000, 5000]);
  const z = x.add(y);
  return z;
});

console.log('Пиковая память:', profile.peakBytes);
profile.result.dispose();

Анализ производительности операций

Профилирование позволяет выявлять узкие места:

  • Каждое ядро (kernel) хранит название операции, входные и выходные формы тензоров, а также время выполнения.
  • Сравнение нескольких профилей помогает оптимизировать последовательность операций и перераспределить ресурсы.
const profile = await tf.profile(() => {
  const a = tf.randomNormal([2000, 2000]);
  const b = tf.randomNormal([2000, 2000]);
  const c = a.matMul(b);
  const d = tf.relu(c);
  return d;
});

profile.kernels.forEach(kernel => {
  console.log(`Операция: ${kernel.name}, Время: ${kernel.kernelTimeMs} мс`);
});
profile.result.dispose();

Профилирование асинхронных операций

TensorFlow.js поддерживает асинхронные операции, такие как tf.data и вызовы к WebGL. tf.profile корректно работает с промисами и асинхронными функциями, собирая данные после завершения всех вычислений внутри колбэка.

const profile = await tf.profile(async () => {
  const data = tf.randomUniform([1000, 1000]);
  const result = await data.square().array();
  return result;
});

console.log('Количество новых тензоров:', profile.newTensors);

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

  • Использовать tf.profile для больших моделей — особенно полезно при обучении нейросетей, чтобы избежать неожиданного переполнения памяти.
  • Освобождать результат профилирования через dispose(), чтобы не создавать дополнительную нагрузку на память.
  • Сохранять профили для сравнительного анализа, например, при изменении структуры модели или размера батча.
  • Анализировать ядра — часто узким местом являются операции матричных умножений или сверток, что важно для оптимизации WebGL и Node.js backend.

Совмещение с tf.engine

tf.engine() предоставляет низкоуровневую информацию о состоянии тензоров. В сочетании с tf.profile можно отслеживать:

  • общее количество тензоров в памяти,
  • потребление памяти по типам данных,
  • детальные сведения о каждом ядре.

Пример интеграции:

const profile = await tf.profile(() => {
  const x = tf.ones([3000, 3000]);
  const y = tf.ones([3000, 3000]);
  return x.matMul(y);
});

console.log('Всего тензоров в движке:', tf.engine().memory().numTensors);
profile.result.dispose();

Использование tf.profile и tf.engine().memory() в комбинации позволяет полностью контролировать потребление ресурсов и оптимизировать вычисления под любые условия, будь то браузер или серверная среда Node.js.

Оглавление