How can I get the symbolic gradient [Tensorflow 2.x]-CodePudding

I want to get the symbolic expression for gradient estimation. When I see the output it's quite difficult to understand what's going on.

import tensorflow as tf
@tf.function
def f_k(input_dat):
    y = tf.matmul(tf.sin(input_dat[0]), input_dat[1])
    grads = tf.gradients([y], input_dat)
    # grads = tape.gradient([y], input_dat)
    tf.print('tf >>', grads)
    print('print >>', grads)
    return y, grads


a = tf.Variable([[1., 3.0], [2., 6.0]])
b = tf.Variable([[1.], [2.]])
input_data = [a, b]
y, z = f_k(input_data)
print(y, z)

Output: inside the function

print >> [<tf.Tensor 'gradients/Sin_grad/mul:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'gradients/MatMul_grad/MatMul_1:0' shape=(2, 1) dtype=float32>]
tf >> [[[0.540302277 -1.979985]
 [-0.416146845 1.92034054]], [[1.75076842]
 [-0.138295487]]

As the output, I want which is shown with print:

[<tf.Tensor 'gradients/Sin_grad/mul:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'gradients/MatMul_grad/MatMul_1:0' shape=(2, 1) dtype=float32>]

However, the function always returns the numerical result. Could someone help me to get this symbolic representation of the gradient?

CodePudding user response：

The symbolic representation you want will only work in graph mode. Outside of graph mode, eager execution is enabled by default. What you can do is create a new function to print the values and wrap it with the @tf.function decorator like you are already doing for f_k:

import tensorflow as tf

@tf.function
def f_k(input_dat):
    y = tf.matmul(tf.sin(input_dat[0]), input_dat[1])
    grads = tf.gradients([y], input_dat)
    # grads = tape.gradient([y], input_dat)
    tf.print('tf >>', grads)
    print('print >>', grads)
    return y, grads

a = tf.Variable([[1., 3.0], [2., 6.0]])
b = tf.Variable([[1.], [2.]])
input_data = [a, b]
y, z = f_k(input_data)

@tf.function
def print_symbolic(y, z):
  print(y,z)
  return y, z
y, z = print_symbolic(y, z)

print >> [<tf.Tensor 'gradients/Sin_grad/mul:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'gradients/MatMul_grad/MatMul_1:0' shape=(2, 1) dtype=float32>]
tf >> [[[0.540302277 -1.979985]
 [-0.416146845 1.92034054]], [[1.75076842]
 [-0.138295487]]]
Tensor("y:0", shape=(2, 1), dtype=float32) [<tf.Tensor 'z:0' shape=(2, 2) dtype=float32>, <tf.Tensor 'z_1:0' shape=(2, 1) dtype=float32>]

You could also just access the tensors of your graph:

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

Tensor("input_dat:0", shape=(), dtype=resource)
Tensor("input_dat_1:0", shape=(), dtype=resource)
Tensor("Sin/ReadVariableOp:0", shape=(2, 2), dtype=float32)
Tensor("Sin:0", shape=(2, 2), dtype=float32)
Tensor("MatMul/ReadVariableOp:0", shape=(2, 1), dtype=float32)
Tensor("MatMul:0", shape=(2, 1), dtype=float32)
Tensor("gradients/Shape:0", shape=(2,), dtype=int32)
Tensor("gradients/grad_ys_0/Const:0", shape=(), dtype=float32)
Tensor("gradients/grad_ys_0:0", shape=(2, 1), dtype=float32)
Tensor("gradients/MatMul_grad/MatMul:0", shape=(2, 2), dtype=float32)
Tensor("gradients/MatMul_grad/MatMul_1:0", shape=(2, 1), dtype=float32)
Tensor("gradients/Sin_grad/Cos:0", shape=(2, 2), dtype=float32)
Tensor("gradients/Sin_grad/mul:0", shape=(2, 2), dtype=float32)
Tensor("StringFormat:0", shape=(), dtype=string)
Tensor("Identity:0", shape=(2, 1), dtype=float32)
Tensor("Identity_1:0", shape=(2, 2), dtype=float32)
Tensor("Identity_2:0", shape=(2, 1), dtype=float32)

Check the docs for more information.