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)

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

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 может быть использован только для вычисления одного набора градиентов (или якобианов).

Помимо тензоров, gradient также поддерживает 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
)

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