@tensorflow/tfjs-node и tfjs-node-gpu

Библиотека TensorFlow.js позволяет выполнять вычисления машинного обучения на платформе Node.js, используя как центральный процессор (CPU), так и графический процессор (GPU). Для Node.js существуют два основных пакета:

  • @tensorflow/tfjs-node – версия для CPU.
  • @tensorflow/tfjs-node-gpu – версия с поддержкой GPU через CUDA и cuDNN.

Установка осуществляется через npm:

npm install @tensorflow/tfjs-node

или для GPU:

npm install @tensorflow/tfjs-node-gpu

После установки подключение выглядит следующим образом:

const tf = require('@tensorflow/tfjs-node'); // или '@tensorflow/tfjs-node-gpu'

Для GPU необходимо наличие совместимого драйвера NVIDIA и установленных библиотек CUDA Toolkit и cuDNN соответствующих версий. Проверка доступности GPU выполняется через метод:

console.log(tf.getBackend());

В Node.js бэкенд по умолчанию tensorflow, что позволяет использовать нативные возможности вычислений, существенно ускоряя операции с тензорами по сравнению с браузерной версией.

Основные типы данных и операции

Тензор (tf.Tensor) является базовым объектом для работы с данными. Создание тензора возможно несколькими способами:

// Тензор из массива
const a = tf.tensor([1, 2, 3, 4]);

// Тензор с заданной формой и значениями по умолчанию
const b = tf.zeros([2, 2]);
const c = tf.ones([3, 3]);

// Тензор случайных чисел
const d = tf.randomNormal([2, 3], 0, 1);

Ключевые операции с тензорами:

  • Арифметические: add, sub, mul, div, pow.
  • Линейная алгебра: matMul, transpose, inverse.
  • Статистические: mean, sum, min, max.
const x = tf.tensor([1, 2, 3]);
const y = tf.tensor([4, 5, 6]);
const sum = x.add(y);
sum.print(); // [5, 7, 9]

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

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

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

Последовательная модель (Sequential)

const model = tf.sequential();

model.add(tf.layers.dense({ units: 32, inputShape: [10], activation: 'relu' }));
model.add(tf.layers.dense({ units: 1, activation: 'linear' }));

model.compile({ optimizer: 'sgd', loss: 'meanSquaredError' });
  • dense – полносвязный слой.
  • inputShape – размерность входного тензора.
  • activation – функция активации.

Обучение модели осуществляется методом fit:

const xs = tf.randomNormal([100, 10]);
const ys = tf.randomNormal([100, 1]);

await model.fit(xs, ys, {
  epochs: 50,
  batchSize: 10,
});

Загрузка и сохранение моделей

// Сохранение модели на диск
await model.save('file://./my-model');

// Загрузка модели
const loadedModel = await tf.loadLayersModel('file://./my-model/model.json');

Сохранение в формате file:// доступно только в Node.js и позволяет эффективно управлять большими проектами.

Работа с GPU

Пакет @tensorflow/tfjs-node-gpu автоматически использует все доступные GPU-устройства. Проверка доступного GPU:

tf.backend().getNumOfPhysicalDevices().then(console.log);

Особенности работы с GPU:

  • Операции матричных умножений и свёртки (convolution) ускоряются в десятки раз.
  • Необходимо следить за использованием видеопамяти, особенно при больших батчах.
  • Переключение между CPU и GPU возможно через метод tf.setBackend('tensorflow') (для Node.js backend автоматически выбирает доступное устройство).

Для эффективной работы рекомендуется использовать:

  • Минимизацию копирования данных между CPU и GPU.
  • Оптимизацию батчей данных.
  • Контроль асинхронного выполнения через tf.tidy для автоматического освобождения памяти.

Асинхронные операции и управление памятью

Все операции TensorFlow.js в Node.js могут быть асинхронными, особенно при использовании GPU. Для предотвращения утечек памяти следует оборачивать вычисления в tf.tidy:

const result = tf.tidy(() => {
  const a = tf.randomNormal([1000, 1000]);
  const b = tf.randomNormal([1000, 1000]);
  return a.matMul(b);
});

После выхода из блока tidy все временные тензоры будут автоматически освобождены, кроме возвращаемого значения.

Для мониторинга используемой памяти:

console.log(tf.memory());

Это возвращает объект с количеством тензоров и объемом занятой памяти, что позволяет эффективно отлаживать проекты с большим числом операций.

Интеграция с другими библиотеками Node.js

TensorFlow.js поддерживает работу с внешними данными:

  • CSV: tf.data.csv('file://data.csv')
  • JSON: fs.readFileSync + tf.tensor
  • Изображения: @tensorflow/tfjs-node содержит модуль node.decodeImage для работы с JPEG и PNG.

Пример загрузки изображения и преобразования в тензор:

const fs = require('fs');
const imageBuffer = fs.readFileSync('./image.jpg');
const imageTensor = tf.node.decodeImage(imageBuffer);

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

Профилирование и оптимизация

TensorFlow.js для Node.js предоставляет инструменты для профилирования:

tf.profile(() => {
  const a = tf.randomNormal([1000, 1000]);
  const b = tf.randomNormal([1000, 1000]);
  return a.matMul(b);
}).then(profile => console.log(profile));

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

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

  • Минимизацию промежуточных тензоров.
  • Выбор оптимального размера батча.
  • Использование асинхронных методов для параллельной обработки данных.
Оглавление