Loop inside tf.function, impossible to compute gradients

user_2309005 0
Viewed 451

I'm trying to compute gradients with respect to some variables defined inside a loop in a tf.function, however I always get a None result. Here is a basic example replicating the problem:

@tf.function
def problem():
  test = tf.constant(1.0)

  with tf.GradientTape() as tape:     
    for i in tf.range(5):
      test1 = test
      tape.watch(test1)
      test2 = test1

  grad = tape.gradient(test2, test1)

  return grad

print(problem())    #None

Of course in this particular case I don't even need the loop. However in a more general situation I would like to store the test1 variable (and possibly others) inside a TensorArray (or similar structure) during the loop, and then compute gradients with respect to those. Is this possible?

tensorflow2.0 tensorflow
2 Answers

Using .range() might work but I think when you are writing a tf operation, having np calls might prevent the operation from running on the GPU. However the problem was raised before, check here.

There is two possible workarounds I can share. One of them is to use tf. range outside the gradient tape.

@tf.function
def problem(step):
    test = tf.constant(1.0)
    for i in tf.range(step):
        with tf.GradientTape() as tape:    
            tape.watch(test)
            test1 = test
            test2 = test1
        grad = tape.gradient(test2, test1)
    return grad

print(problem(5)) # 1

or run without @tf.function.

def problem(step):
    test = tf.constant(1.0)
    with tf.GradientTape() as tape:   
        for _ in tf.range(step):
            tape.watch(test)
            test1 = test
            test2 = test1
    grad = tape.gradient(test2, test1)
    return grad

print(problem(5))  # 1

Or, you can look into tf.while_loop.

Update

.range is good to use. It won't create any major bottleneck.

user_92157800

I can't fully explain why this works, but you can replace tf.range with range to get a desired result. Note that you could keep your original code but remove the tf.function decorator, and you would get the same result.

import tensorflow as tf

@tf.function
def problem():
  test = tf.constant(1.0)
  with tf.GradientTape() as tape:     
    for i in range(5):
      test1 = test
      tape.watch(test1)
      test2 = test1

  grad = tape.gradient(test2, test1)

  return grad

print(problem())  # tf.Tensor(1.0, shape=(), dtype=float32)
user_56660870
Related