tf.GradientTape 如何记录 with 语句内部的操作？

Question

我不明白 tf.GradientTape 如何在“with”语句（以下操作）中记录像y=x**2这样的操作。

x = tf.Variable(3.0)

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

可以使用什么 Python 语法来实现此行为？

Answer 1

编辑：

根据 GitHub 源代码， GradientTape ，在Line 897 ：

 @tf_contextlib.contextmanager
  def _ensure_recording(self):
    """Ensures that this tape is recording."""
    if not self._recording:
      try:
        self._push_tape()
        yield
      finally:
        self._pop_tape()
    else:
      yield

如果您不知道，只要with关键字， contextmanager就会触发。 它告诉我们它开始跟踪磁带。

self._pop_tape()在Line 891 ：

def _pop_tape(self):
    if not self._recording:
      raise ValueError("Tape is not recording.")
    tape.pop_tape(self._tape)
    self._recording = False

self._push_tape()在Line 878 ：

def _push_tape(self):
    """Pushes a new tape onto the tape stack."""
    if self._recording:
      raise ValueError("Tape is still recording, This can happen if you try to "
                       "re-enter an already-active tape.")
    if self._tape is None:
      self._tape = tape.push_new_tape(
          persistent=self._persistent,
          watch_accessed_variables=self._watch_accessed_variables)
    else:
      tape.push_tape(self._tape)
    self._recording = True

在这里，您可以注意到tape.push_new_tape在这里被访问，可以在Line 43的源代码中找到：

def push_new_tape(persistent=False, watch_accessed_variables=True):
  """Pushes a new tape onto the tape stack."""
  tape = pywrap_tfe.TFE_Py_TapeSetNew(persistent, watch_accessed_variables)
  return Tape(tape)

在这里，您可以在Line 31上方看到Tape class。

class Tape(object):
  """Represents a gradient propagation trace."""

  __slots__ = ["_tape"]

  def __init__(self, tape):
    self._tape = tape

  def watched_variables(self):
    return pywrap_tfe.TFE_Py_TapeWatchedVariables(self._tape)

此外，我试图跟踪pywrap_tfe.TFE_Py_TapeSetNew但无法在文件的源代码中找到它。

原始答案：

GradientTape的文档指出：

默认情况下，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