python - 如何使用 tensorflow gradienttape 找到分析梯度

Question

假设我们有一些函数 y=x^2

然后我们可以使用梯度带自动为我们计算梯度（当我们提供一些 x 的值给 tensorflow

x = tf.Variable(3.0)

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

无论如何我可以找出 tensorflow 对我的输入做了什么？例如在这种情况下很容易找出 dy/dx=2x，这是否意味着 tensorflow 会将 2 乘以我的 x 输入值，然后返回 6（即 3*2）？

我有一个非常复杂的函数，我不知道如何区分，所以我想从 tensorflow gradienttape 中找到见解，看看 tensorflow 如何使用我的 x 输入计算导数。

Answer 1

一种可能的选择是使用tensorboard并打印以下操作tf.Graph：

import tensorflow as tf

x = tf.Variable(3.0)
@tf.function
def f(x):
  with tf.GradientTape() as tape:
    y = x**2
  return tape.gradient(y, x)

print(*[tensor for op in f.get_concrete_function(x).graph.get_operations() for tensor in op.values()], sep="\n")

logdir = 'logs/func/'
writer = tf.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True, profiler=True)
z = f(x)
with writer.as_default():
  tf.summary.trace_export(
      name="func_trace",
      step=0,
      profiler_outdir=logdir)

Tensor("x:0", shape=(), dtype=resource)
Tensor("ReadVariableOp:0", shape=(), dtype=float32)
Tensor("pow/y:0", shape=(), dtype=float32)
Tensor("pow:0", shape=(), dtype=float32)
Tensor("ones:0", shape=(), dtype=float32)
Tensor("gradient_tape/pow/mul:0", shape=(), dtype=float32)
Tensor("gradient_tape/pow/sub/y:0", shape=(), dtype=float32)
Tensor("gradient_tape/pow/sub:0", shape=(), dtype=float32)
Tensor("gradient_tape/pow/Pow:0", shape=(), dtype=float32)
Tensor("gradient_tape/pow/mul_1:0", shape=(), dtype=float32)
Tensor("Identity:0", shape=(), dtype=float32)

在终端中打开张量板：

%load_ext tensorboard
%tensorboard --logdir logs/func