tf.customGrad

TensorFlow.js предоставляет высокоуровневые абстракции для работы с нейронными сетями, однако иногда требуется полный контроль над процессом обратного распространения ошибки. Для этого используется функция tf.customGrad, которая позволяет определить собственный градиент для произвольной операции.

Основы tf.customGrad

tf.customGrad принимает один аргумент — функцию, которая возвращает объект с двумя свойствами:

  1. value — результат прямого вычисления операции.
  2. gradFunc — функция, вычисляющая градиент по отношению к входным тензорам.

Простейший синтаксис выглядит так:

const customOp = tf.customGrad((x, save) => {
  save([x]);  // Сохраняем тензор для использования в gradFunc
  const value = x.square();
  const gradFunc = (dy, saved) => {
    const [x] = saved;
    return dy.mul(x.mul(2));  // Производная x^2 равна 2*x
  };
  return { value, gradFunc };
});

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

Аргументы gradFunc

Функция gradFunc принимает два параметра:

  1. dy — градиент по отношению к результату функции (value).
  2. saved — массив тензоров, сохраненных с помощью save.

Возвращаемое значение gradFunc должно соответствовать количеству входных тензоров, переданных в customGrad. Для нескольких входов градиенты возвращаются в виде массива:

const customAdd = tf.customGrad((a, b, save) => {
  save([a, b]);
  const value = a.add(b);
  const gradFunc = (dy, saved) => {
    return [dy, dy];  // Производная суммы по каждому аргументу равна 1
  };
  return { value, gradFunc };
});

Сохранение тензоров для градиента

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

const customExp = tf.customGrad((x, save) => {
  const y = x.exp();
  save([y]);  // Сохраняем результат экспоненты
  const gradFunc = (dy, saved) => {
    const [y] = saved;
    return dy.mul(y);  // Производная e^x равна e^x
  };
  return { value: y, gradFunc };
});

Применение tf.customGrad в обучении

Пользовательские градиенты особенно полезны при:

  • Оптимизации нестандартных функций потерь.
  • Реализации операций с нестандартным поведением градиента (например, аппроксимация функции с разрывами).
  • Экспериментировании с новыми архитектурами нейросетей, где стандартные градиенты могут быть недостаточно эффективными.

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

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

for (let i = 0; i < 100; i++) {
  optimizer.minimize(() => {
    const x = tf.tensor1d([1, 2, 3]);
    const y = customOp(x);
    return y.sum();  // Минимизируем сумму квадратов
  });
}

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

Советы по эффективности

  • Сохранять только необходимые тензоры с помощью save. Избыточное сохранение приводит к увеличению потребления памяти.
  • Использовать tf.tidy внутри customGrad при необходимости управлять временными тензорами.
  • Проверять вычисляемые градиенты через tf.grad или tf.valueAndGrad для отладки сложных операций.

Примеры нестандартных градиентов

Градиент с пороговой функцией (clipping)

const customClip = tf.customGrad((x, save) => {
  const y = x.relu();  // ReLU
  save([x]);
  const gradFunc = (dy, saved) => {
    const [x] = saved;
    return dy.mul(x.greater(0).toFloat());  // Градиент ReLU
  };
  return { value: y, gradFunc };
});

Градиент с масштабированием

const scaledGrad = tf.customGrad((x, save) => {
  const y = x.square();
  save([x]);
  const gradFunc = (dy, saved) => {
    const [x] = saved;
    return dy.mul(x.mul(2).mul(0.5));  // Масштабируем производную
  };
  return { value: y, gradFunc };
});

Эти примеры демонстрируют гибкость tf.customGrad для создания операций с нестандартным поведением обратного распространения.

Взаимодействие с другими функциями TensorFlow.js

tf.customGrad полностью совместим с автоматическим дифференцированием. Это означает, что созданные с его помощью операции можно комбинировать с обычными операциями tf.add, tf.mul, tf.matMul и другими, и TensorFlow.js корректно вычислит градиенты по всему графу.

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

Оглавление