Реализация автоматического дифференцирования. Объект GradientTape

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

Начнем с довольно простого примера. Предположим, нам нужно вычислить производную функции  в точке . Для этого мы создадим переменную со значением -2.0 (указание вещественной величины здесь обязательно) и выполним прямой проход по графу вычислений для заданной функции, используя специальный объект GradientTape:

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
import tensorflow as tf
 
x = tf.Variable(-2.0)
 
with tf.GradientTape() as tape:
    y = x ** 2
 
df = tape.gradient(y, x)
print(df)

Смотрите, когда мы запускаем менеджер контекста (оператор with) для объекта GradientTape(), то дляфункции, записанной в виде y = x ** 2, сохраняются все необходимые вычисления (для дальнейшего нахождения производных) в объекте tape.После этого вызывается метод gradient(), который реализует вычисление частных производных методом обратного распространения (приобратом проходе по графу). (Если вы не знаете, как работает вычислительный граф, то смотрите первое занятие этой темы). Здесь первым параметром мы указываем саму функцию y, а вторым – аргументы, от которых вычисляются производные. В данном случае указан один аргумент x. Результат вычисления будет представлен следующим образом:

tf.Tensor(-4.0, shape=(), dtype=float32)

То есть, мы получили значение производной в точке :

Может показаться, что городить такой огород для вычисления производных – это верх расточительства ресурсов компьютера. В конце концов, ее и подобные ей производные можно посчитать вручную и просто реализовать в виде функции. Все верно, для таких простых ситуаций именно так и следует поступать. Однако, в задачах машинного обучения часто имеют дело с функциями, зависящими от огромного числа аргументов (тысячи и миллионы). В этих условиях напрямую производить вычисления через якобиан – неэффективное решение, а иногда и непреодолимое из-за отсутствия достаточного объема памяти. А вот с помощью вычислительного графа такие задачи решаются и относительно просто.

Как пример, давайте возьмем функцию от вектора параметров w,b и x. Вид функции

напоминает вычисления на входах нейронов с весовыми коэффициентами w и смещениями b. Затем, мы на ее основе определим следующую функцию вида:

для которой и вычислим частные производные по вектору w и b:

w = tf.Variable(tf.random.normal((3, 2)))
b = tf.Variable(tf.zeros(2, dtype=tf.float32))
x = tf.Variable([[-2.0, 1.0, 3.0]])
 
with tf.GradientTape() as tape:
    y = x @ w + b
    loss = tf.reduce_mean(y ** 2)
 
df = tape.gradient(loss, [w, b])
print(df[0], df[1], sep="\n")

Смотрите, как это легко реализуется на Tensorflow с помощью процедуры автоматического дифференцирования. В методе gradient() мы первым параметром указываем функцию loss, а вторым – список аргументов, от которых вычислим частные производные в точках w,b и при заданном значении. Затем, эти производные можно использовать для нахождения экстремума функции loss (в данном случае ее минимума).

Но как объект GradientTape() «узнает», какие промежуточные вычисления ему нужно отслеживать и сохранять в переменной tape?В Tensorflow это сделано относительно просто. Если функция зависит от каких-либо объектов tf.Variable(), то они и отслеживаются для последующей возможности вычисления градиентов. Давайте на простом примере в этом убедимся. Возьмем функцию, зависящую от двух параметров:

x = tf.Variable(0, dtype=tf.float32)
b = tf.constant(1.5)
 
with tf.GradientTape() as tape:
    f = (x + b) ** 2 + 2 * b
 
df = tape.gradient(f, [x, b])
print(df[0], df[1], sep="\n")

После ее выполнения в консоли увидим:

tf.Tensor(3.0, shape=(), dtype=float32)
None

Это говорит о том, что по переменной x производная была вычислена, а по константе b нет. Но, если константу преобразовать в переменную:

b = tf.Variable(1.5)

то мы увидим обе производные.

Также в Tensorflow можно запретить отслеживать какие-либо переменные, указав у них параметр trainable=False:

x = tf.Variable(0, dtype=tf.float32, trainable=False)

Тогда в объекте tapeбудут отсутствовать данные для вычисления производной по x. Или, вот такой не очень очевидный случай:

b = tf.Variable(1.5) + 1.0

Здесь переменная b после операции сложения превращается в объект tf.Tensor и перестает отслеживаться объектом GradientTape().

Все эти моменты следует учитывать при разработке программ с автоматическим дифференцированием. Но зачем в Tensorflow все так реализовано? Почему бы не отслеживать все переменные, от которых зависит целевая функция?Дело в том, что промежуточные вычисления могут занимать большой объем памяти. Представьте НС с десятками тысяч нейронов и миллионами весовых коэффициентов. Для каждого из параметров будут сохраняться значения вычислений и это не малый объем. Поэтому, при реализации GradientTape() можно явно указать переменные, которые нужно отслеживать, чтобы не расходовать зря память устройства и не делать лишних промежуточных вычислений. Мало того, с помощью параметра watch_accessed_variables=False можно полностью отключать отслеживание переменных и тогда GradientTape() не будет производить никаких дополнительных вычислений:

x = tf.Variable(0, dtype=tf.float32)
b = tf.Variable(1.5)
 
with tf.GradientTape(watch_accessed_variables=False) as tape:
    f = (x + b) ** 2 + 2 * b
 
df = tape.gradient(f, [x, b])
print(df[0], df[1], sep="\n")

Увидим значения Noneдля обеих производных. Получается, что в таком режиме мы вовсе не можем вычислять производные? Не совсем. В самом менеджере контекста можно явно указать те переменные, которые следует наблюдать. Это делается с помощью метода watch():

with tf.GradientTape(watch_accessed_variables=False) as tape:
    tape.watch(x)
    f = (x + b) ** 2 + 2 * b

Теперь производная по x может быть вычислена. А если указать обе переменные:

tape.watch([x, b])

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

Также отслеживание происходит и для всех промежуточных переменных, которые связаны с наблюдаемой переменной. В примере ниже наблюдается переменная x, через нее вычисляется промежуточная переменная y, а затем, функция f:

with tf.GradientTape(watch_accessed_variables=False) as tape:
    tape.watch(x)
    y = 2 * x
    f = y * y
 
df = tape.gradient(f, y)
print(df)

Соответственно, мы можем вычислить градиент как поy, так и поx. Однако, обратите внимание, если мы попытаемся вычислить градиенты последовательным вызовом методаgradient():

df = tape.gradient(f, y)
df_dx = tape.gradient(f, x)

То при повторном вызове получим ошибку. Это связано с тем, что метод gradient() автоматически высвобождает все ресурсы, связанные с промежуточными вычислениями. Чтобы этого не происходило, мы можем указать параметр persistent=True при создании объекта GradientTape():

with tf.GradientTape(watch_accessed_variables=False, persistent=True) as tape:

Теперь производные будут успешно посчитаны. Но, освобождение ресурсов теперь на наших плечах (или плечах сборщика мусора). Здесь лучше явно записать команду:

del tape

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

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

x = tf.Variable(1.0)
 
with tf.GradientTape() as tape:
    y = [2.0, 3.0] * x ** 2
 
df = tape.gradient(y, x)
print(df)

То есть, здесь, фактически происходит суммирование градиентов от каждого выходного значения.

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

x = tf.Variable([1.0, 2.0])
 
with tf.GradientTape() as tape:
    y = tf.reduce_sum([2.0, 3.0] * x ** 2)
 
df = tape.gradient(y, x)
print(df)

Так как, проходя по графу в обратном направлении, как раз получаем две величины для каждого входа.

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

x = tf.Variable(1.0)
 
with tf.GradientTape() as tape:
    if x < 2.0:
        y = tf.reduce_sum([2.0, 3.0] * x ** 2)
    else:
        y = x ** 2
 
df = tape.gradient(y, x)
print(df)

Это возможно благодаря вычислению градиентов в конкретной точке (при заданном значении x). Если x=1, как в нашем примере, то по условию будет выбрана первая функция и для нее посчитаны промежуточные значения. Иначе, будет взята вторая функция. То есть, сам объект GradientTape(), конечно же, работает с конкретной функцией и не вычисляет производные от условных операторов.

Особенности вычисления градиентов

В заключение этого занятия отмечу некоторые частые ошибки, совершаемые новичками, при использовании объекта GradientTape(). Первая связана с неверным определением промежуточной функции. Например, мы вычисляем градиент, следующим образом:

x = tf.Variable(1.0)
y = 2 * x + 1
 
with tf.GradientTape() as tape:
    z = y ** 2
 
df = tape.gradient(z, x)

Казалось бы, здесь y зависит от x, а zот y, поэтому можно вычислить производную zпо x. Но на выходе получим значение None. Это из-за того, что промежуточные значения для функции y, зависящей от x, не были записаны в объект tape, так как она определена вне контекста выполнения объекта GradientTape(). Поправить это очень просто. Поместим строчку с функцией y внутрь контекста:

with tf.GradientTape() as tape:
    y = 2 * x + 1
    z = y ** 2

И тогда сможем получить требуемый результат.

Следующая частая ошибка, когда в процессе вычисления градиента случайно заменяем исходную переменную x на тензор (константу):

x = tf.Variable(1.0)
 
for n in range(2):
    with tf.GradientTape() as tape:
        y = x ** 2 + 2 * x
 
    df = tape.gradient(y, x)
    print(df)
 
    x = x + 1 # так делать нельзя

Здесь первый раз градиент будет вычислен успешно, а второе значение будет None, так как после операции x = x + 1 переменная x превращается в тензор, который не отслеживается в GradientTape(). Правильнее здесь было бы использовать специальный метод assing_add(), следующим образом:

x.assign_add(1.0)

И тогда оба значения будут посчитаны корректно.

Следующая ошибка связана с вычислением значений не методами Tensorflow, а другими способами. Например, так:

x = tf.Variable(1.0)
 
with tf.GradientTape() as tape:
    y = tf.constant(2.0) + np.square(x)
 
df = tape.gradient(y, x)
print(df)

Здесь производится возведение переменной xв квадрат средствами пакета NumPy, в результате, объект GradientTape() не имеет возможности отслеживать поведение x, так как мы как бы выходим из Tensorflow и производная уже не может быть посчитана. А вот если вместо NumPy записать просто x * x:

y = tf.constant(2.0) + x * x

то получим нужный результат.

Следующая частая ошибка – это определение переменной с целым типом данных:

x = tf.Variable(1)
 
with tf.GradientTape() as tape:
    y = x * x
 
df = tape.gradient(y, x)

В вычислительном графе также начинают фигурировать целочисленные величины и как результат потеря точности промежуточных величин. В этих условиях Tensorflow не производит вычисления градиента.

Наконец последнее, что я отмечу на этом занятии – это необходимость прописывать формулы для целевой функции в явном виде. Представим, что у нас имеются две переменные x и w и в контексте GradientTape() производится их суммирование методом assign_add():

x = tf.Variable(1.0)
w = tf.Variable(2.0)
 
with tf.GradientTape() as tape:
    w.assign_add(x)
    y = w ** 2
 
df = tape.gradient(y, x)
print(df)

В результате объект GradientTape() теряет явную связь переменной w с x и, как следствие, нельзя вычислить производную yпо x. Правильно здесь было бы записать так:

with tf.GradientTape() as tape:
    w = w + x
    y = w ** 2

Надеюсь, из материала этого занятия вы стали лучше понимать порядок вычисления градиентов в Tensorflow с помощью объекта GradientTape().

Видео по теме