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

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

Примечание: По умолчанию реализация jacobian использует параллельный цикл (pfor), который неявно создаёт tf.function при каждом вызове jacobian. Для повышения производительности и предотвращения повторной компиляции и преобразований векторизации при каждом вызове поместите код 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
)

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