tf.GradientTape

Запись операций для автоматического дифференцирования.

tf.GradientTape(
    persistent=False, watch_accessed_variables=True
)

Используется в блокнотах

Операции записываются, если они выполняются в рамках этого контекстного менеджера, и по меньшей мере один из их входных данных «отслеживается».

Обучаемые переменные (созданные с помощью tf.Variable или tf.compat.v1.get_variable, где trainable=True — значение по умолчанию в обоих случаях) автоматически отслеживаются. Тензоры можно вручную отслеживать, вызвав метод watch в этом контекстном менеджере.

Например, рассмотрим функцию y = x * x. Градиент в x = 3.0 может быть вычислен как:

x = tf.constant(3.0)
with tf.GradientTape() as g:
  g.watch(x)
  y = x * x
dy_dx = g.gradient(y, x)
print(dy_dx)
tf.Tensor(6.0, shape=(), dtype=float32)

GradientTape можно вкладывать для вычисления производных высших порядков. Например,

x = tf.constant(5.0)
with tf.GradientTape() as g:
  g.watch(x)
  with tf.GradientTape() as gg:
    gg.watch(x)
    y = x * x
  dy_dx = gg.gradient(y, x)  # dy_dx = 2 * x
d2y_dx2 = g.gradient(dy_dx, x)  # d2y_dx2 = 2
print(dy_dx)
tf.Tensor(10.0, shape=(), dtype=float32)
print(d2y_dx2)
tf.Tensor(2.0, shape=(), dtype=float32)

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

x = tf.constant(3.0)
with tf.GradientTape(persistent=True) as g:
  g.watch(x)
  y = x * x
  z = y * y
dz_dx = g.gradient(z, x)  # (4*x^3 at x = 3)
print(dz_dx)
tf.Tensor(108.0, shape=(), dtype=float32)
dy_dx = g.gradient(y, x)
print(dy_dx)
tf.Tensor(6.0, shape=(), dtype=float32)

По умолчанию GradientTape автоматически отслеживает все обучаемые переменные, к которым обращаются внутри контекста. Если вы хотите иметь более точный контроль над отслеживаемыми переменными, вы можете отключить автоматическое отслеживание, передав watch_accessed_variables=False конструктору ленты:

x = tf.Variable(2.0)
w = tf.Variable(5.0)
with tf.GradientTape(
    watch_accessed_variables=False, persistent=True) as tape:
  tape.watch(x)
  y = x ** 2  # Gradients will be available for `x`.
  z = w ** 3  # No gradients will be available as `w` isn't being watched.
dy_dx = tape.gradient(y, x)
print(dy_dx)
tf.Tensor(4.0, shape=(), dtype=float32)
# No gradients will be available as `w` isn't being watched.
dz_dw = tape.gradient(z, w)
print(dz_dw)
None

Обратите внимание, что при использовании моделей вы должны убедиться, что ваши переменные существуют при использовании watch_accessed_variables=False. В противном случае очень легко сделать так, чтобы у вас не было градиентов в первом итерации:

a = tf.keras.layers.Dense(32)
b = tf.keras.layers.Dense(32)

with tf.GradientTape(watch_accessed_variables=False) as tape:
  tape.watch(a.variables)  # Since `a.build` has not been called at this point
                           # `a.variables` will return an empty list and the
                           # tape will not be watching anything.
  result = b(a(inputs))
  tape.gradient(result, a.variables)  # The result of this computation will be
                                      # a list of `None`s since a's variables
                                      # are not being watched.

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

Методы

batch_jacobian

Просмотреть исходный код

batch_jacobian(
    target,
    source,
    unconnected_gradients=tf.UnconnectedGradients.NONE,
    parallel_iterations=None,
    experimental_use_pfor=True
)

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

См. статью Википедии для определения якобиана. Эта функция по сути является эффективной реализацией следующего:

tf.stack([self.jacobian(y[i], x[i]) for i in range(x.shape[0])]).

Обратите внимание, что по сравнению с GradientTape.jacobian, который вычисляет градиент каждого выходного значения по отношению к каждому входному значению, эта функция полезна, когда target[i,...] не зависит от source[j,...] для j != i. Это предположение позволяет более эффективно вычислять, по сравнению с GradientTape.jacobian. Выходные данные, а также промежуточные активации, имеют меньшую размерность и избегают большого количества излишних нулей, что привело бы к вычислению якобиана при предположении независимости.

Примечание: Без установки persistent=True, GradientTape можно использовать только для вычисления одного набора градиентов (или якобианов).

Примечание: По умолчанию реализация batch_jacobian использует параллельное вычисление (pfor), что создает tf.function для каждого вызова batch_jacobian. Для повышения производительности и предотвращения повторной компиляции и переписывания векторизации при каждом вызове, поместите код GradientTape в @tf.function.

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

with tf.GradientTape() as g:
  x = tf.constant([[1., 2.], [3., 4.]], dtype=tf.float32)
  g.watch(x)
  y = x * x
batch_jacobian = g.batch_jacobian(y, x)
# batch_jacobian is [[[2,  0], [0,  4]], [[6,  0], [0,  8]]]

gradient

Просмотреть исходный код

gradient(
    target,
    sources,
    output_gradients=None,
    unconnected_gradients=tf.UnconnectedGradients.NONE
)

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

Примечание: Без установки persistent=True, GradientTape можно использовать только для вычисления одного набора градиентов (или якобианов).

Помимо тензоров, градиент также поддерживает RaggedTensors. Например,

x = tf.ragged.constant([[1.0, 2.0], [3.0]])
with tf.GradientTape() as g:
  g.watch(x)
  y = x * x
g.gradient(y, x)
<tf.RaggedTensor [[2.0, 4.0], [6.0]]>

jacobian

Просмотреть исходный код

jacobian(
    target,
    sources,
    unconnected_gradients=tf.UnconnectedGradients.NONE,
    parallel_iterations=None,
    experimental_use_pfor=True
)

Вычисляет якобиан, используя операции, записанные в контексте этой ленты.

Примечание: Если вы не установили persistent=True, GradientTape можно использовать только для вычисления одного набора градиентов (или якобианов).

Примечание: По умолчанию реализация якобиана использует параллельную обработку (pfor), которая создаёт tf.function для каждого вызова якобиана. Для повышения производительности и для предотвращения повторной компиляции и переписывания векторизации при каждом вызове, заключите код GradientTape в @tf.function.

См.статью Википедии для определения якобиана.

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

with tf.GradientTape() as g:
  x  = tf.constant([1.0, 2.0])
  g.watch(x)
  y = x * x
jacobian = g.jacobian(y, x)
# jacobian value is [[2., 0.], [0., 4.]]

reset

Просмотреть исходный код

reset()

Очищает всю информацию, сохранённую в этой ленте.

Эквивалентно выходу и повторному входу в контекст менеджера ленты с новой лентой. Например, следующие два блока кода эквивалентны:

with tf.GradientTape() as t:
  loss = loss_fn()
with tf.GradientTape() as t:
  loss += other_loss_fn()
t.gradient(loss, ...)  # Only differentiates other_loss_fn, not loss_fn


# The following is equivalent to the above
with tf.GradientTape() as t:
  loss = loss_fn()
  t.reset()
  loss += other_loss_fn()
t.gradient(loss, ...)  # Only differentiates other_loss_fn, not loss_fn

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

with tf.GradientTape() as t:
  loss = ...
  if loss > k:
    t.reset()

stop_recording

Просмотреть исходный код

@tf_contextlib.contextmanager
stop_recording()

Временно останавливает запись операций на этой ленте.

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

Например:

x = tf.constant(4.0)
with tf.GradientTape() as tape:
  with tape.stop_recording():
    y = x ** 2
dy_dx = tape.gradient(y, x)
print(dy_dx)
None

watch

Просмотреть исходный код

watch(
    tensor
)

Обеспечивает, что tensor отслеживается этой лентой.

watched_variables

Просмотреть исходный код

watched_variables()

Возвращает отслеживаемые этой лентой переменные в порядке их создания.

__enter__

Просмотреть исходный код

__enter__()

Входит в контекст, внутри которого операции записываются на этой ленте.

__exit__

Просмотреть исходный код

__exit__(
    typ, value, traceback
)

Выходит из контекста записи, больше никаких операций не отслеживаются.