Why does Tensorflow's automatic differentiation fail when .numpy() is used in the loss function?

user_9897685 0
Viewed 296

I've noticed that Tensorflow's automatic differentiation does not give the same values as finite differences when the loss function converts the input to a numpy array to calculate the output value. Here's a minimum working example of the problem:

import tensorflow as tf
import numpy as np

def lossFn(inputTensor):
    # Input is a rank-2 square tensor
    return tf.linalg.trace(inputTensor @ inputTensor)

def lossFnWithNumpy(inputTensor):
    # Same function, but converts input to a numpy array before performing the norm
    inputArray = inputTensor.numpy()

    return tf.linalg.trace(inputArray @ inputArray)

N = 2
tf.random.set_seed(0)
randomTensor = tf.random.uniform([N, N])

# Prove that the two functions give the same output; evaluates to exactly zero
print(lossFn(randomTensor) - lossFnWithNumpy(randomTensor)) 

theoretical, numerical = tf.test.compute_gradient(lossFn, [randomTensor])
# These two values match
print(theoretical[0])
print(numerical[0])

theoretical, numerical = tf.test.compute_gradient(lossFnWithNumpy, [randomTensor])
# The theoretical value is [0 0 0 0]
print(theoretical[0])
print(numerical[0])

The function tf.test.compute_gradients computes the 'theoretical' gradient using automatic differentiation, and the numerical gradient using finite differences. As the code shows, if I use .numpy() in the loss function the automatic differentiation does not calculate the gradient.

Could anybody explain the reason for this?

tensorflow numpy python
1 Answers

From the guide : Introduction to Gradients and Automatic Differentiation

The tape can't record the gradient path if the calculation exits TensorFlow. For example:

x = tf.Variable([[1.0, 2.0],
                 [3.0, 4.0]], dtype=tf.float32)

with tf.GradientTape() as tape:   
  x2 = x**2
  # This step is calculated with NumPy   
  y = np.mean(x2, axis=0)
  # Like most ops, reduce_mean will cast the NumPy array to a constant tensor 
  # using `tf.convert_to_tensor`. 
  y = tf.reduce_mean(y,axis=0)

print(tape.gradient(y, x))

outputs None

The numpy value will be cast back as a constant tensor in the call to tf.linalg.trace, which Tensorflow cannot compute gradients on.

user_73701530
Related