tf.grad и tf.grads

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

Функция tf.grad

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

Синтаксис:

const gradFunction = tf.grad(f);
const gradient = gradFunction(x);
  • f — функция, возвращающая скалярное значение (тензор с размерностью 0).
  • x — тензор, относительно которого вычисляется градиент.
  • gradient — тензор того же типа и формы, что x, содержащий производные функции f по x.

Пример:

const x = tf.tensor1d([1, 2, 3]);

const f = x => x.square().sum(); // f(x) = x^2 + y^2 + z^2
const gradF = tf.grad(f);

const dx = gradF(x);
dx.print(); // [2, 4, 6]

Пояснение: функция x.square().sum() возвращает скаляр, и tf.grad вычисляет частные производные по каждому элементу тензора x. Результат [2, 4, 6] соответствует ∂f/∂xᵢ для каждого элемента.

Функция tf.grads

Если необходимо вычислить градиенты функции относительно нескольких тензоров одновременно, используется tf.grads. Она работает аналогично tf.grad, но возвращает массив градиентов.

Синтаксис:

const gradsFunction = tf.grads(f);
const gradients = gradsFunction([x1, x2, ...]);
  • f — функция, возвращающая скалярный тензор.
  • [x1, x2, ...] — массив тензоров, по которым требуется дифференцирование.
  • gradients — массив градиентов, каждый элемент соответствует входному тензору.

Пример:

const x = tf.tensor1d([1, 2, 3]);
const y = tf.tensor1d([4, 5, 6]);

const f = (x, y) => x.square().sum().add(y.square().sum());
const gradF = tf.grads(f);

const [dx, dy] = gradF([x, y]);
dx.print(); // [2, 4, 6]
dy.print(); // [8, 10, 12]

Пояснение: tf.grads вычисляет производные функции f одновременно по каждому тензору. Это особенно удобно при обучении моделей с несколькими входами.

tf.variable и градиенты

Чтобы градиенты могли использоваться для обновления параметров модели, тензоры должны быть изменяемыми. Для этого применяется tf.variable:

const w = tf.variable(tf.tensor1d([0.5, -0.5]));

const f = w => w.square().sum();
const gradF = tf.grad(f);

const dw = gradF(w);
w.assign(w.sub(dw.mul(0.1))); // шаг градиентного спуска

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

Использование tf.tidy для управления памятью

Вычисление градиентов создает временные тензоры, которые необходимо освобождать, чтобы избежать утечек памяти. Для этого применяется tf.tidy:

const gradF = tf.grad(x => tf.tidy(() => x.square().sum()));
const dx = gradF(tf.tensor1d([1, 2, 3]));
dx.print();

tf.tidy автоматически очищает промежуточные тензоры, оставляя только возвращаемый результат.

Градиенты и оптимизация

Для оптимизации модели можно напрямую использовать градиенты вместе с встроенными оптимизаторами TensorFlow.js:

const optimizer = tf.train.sgd(0.1);

const w = tf.variable(tf.tensor1d([1, -1]));

optimizer.minimize(() => w.square().sum());
w.print(); // значение после одного шага оптимизации

optimizer.minimize автоматически вычисляет градиенты через tf.grad и обновляет переменные, что упрощает реализацию обучения.

Векторизация градиентов

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

const x = tf.tensor2d([[1, 2], [3, 4]]);

const gradF = tf.grad(x => x.square().sum());
gradF(x).print(); // [[2, 4], [6, 8]]

Ограничения и важные моменты

  1. tf.grad и tf.grads подходят только для функций, возвращающих скалярный тензор. Для работы с функциями, возвращающими тензоры любой формы, используется tf.valueAndGrad или tf.valueAndGrads.
  2. Все вычисления производятся лениво, фактическое значение градиента создается только при вызове .data() или .print().
  3. Для сложных графов с множеством операций рекомендуется использовать tf.tidy, чтобы контролировать память и избегать утечек.

Практическая схема работы

  1. Создание переменных через tf.variable.
  2. Определение функции потерь, которая возвращает скаляр.
  3. Вычисление градиентов с помощью tf.grad или tf.grads.
  4. Обновление переменных вручную или через оптимизатор.
  5. Очистка промежуточных тензоров с помощью tf.tidy.

Эта схема является фундаментальной для построения нейронных сетей и реализации алгоритмов оптимизации в TensorFlow.js.

Оглавление