如何在用@tf.function 装饰的函数内部使用 for 循环来操作和返回 tf.Variable tf.data.Dataset？

Question

I am trying to create a function containing a for loop over a TensorFlow Dataset that assigns a new value to a TensorFlow Variable in each iteration. The Variable should also be returned as output of the function. With eager execution enabled, there are no issues, however, in graph mode, some unexpected things seem to happen. Consider the following simple dummy code:

import tensorflow as tf


class Test(object):
    def __init__(self):
        self.var = tf.Variable(0, trainable=False, dtype=tf.float32)
        self.increment = tf.constant(1, dtype=tf.float32)
        self.dataset = tf.data.Dataset.from_tensor_slices([0, 1, 2])

    @tf.function
    def fn1(self):
        self.var.assign(0)
        for _ in tf.range(3):
            self.var.assign(self.var+self.increment)
            tf.print(self.var)
        tf.print(self.var)
        return self.var

    @tf.function
    def fn2(self):
        self.var.assign(0)
        for _ in self.dataset:
            self.var.assign(self.var+self.increment)
            tf.print(self.var)
        tf.print(self.var)
        return self.var

    @tf.function
    def fn3(self):
        self.var.assign(0)
        y = self.var
        for _ in self.dataset:
            self.var.assign(self.var+self.increment)
            y = self.var
            tf.print(y)
        tf.print(y)
        return y

    @tf.function
    def fn4(self):
        var = 0.0
        for _ in self.dataset:
            var += 1.0
            tf.print(var)
        tf.print(var)
        return var

test.fn1() , test.fn3() and test.fn4() all return the following (desired) output:

1
2
3
3
<tf.Tensor: shape=(), dtype=float32, numpy=3.0>

However, test.fn2() behaves differently:

1
2
3
0
<tf.Tensor: shape=(), dtype=float32, numpy=0.0>

Interestingly, after execution of test.fn2 , test.var does seem to contain the correct value:

<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>

I am not sure why test.fn2 fails. Clearly, it is doing some things correctly (as test.var contains the correct value after execution of the function), but it does not deliver the correct result. Can you help me understand what causes this code to fail?

The behavior described above occurs when using TensorFlow 2.1.0 for Python 3.6 on CentOS 7.

Answer 1

Running this on TensorFlow 2.1.0 reproduces your scenario.

Which prints 1 2 3 0 for test.fn2() , but you should also consider that when you print self.var in test.fn3() it will also show you self.var = 0 during the function call.

Modified fn3( ) :

    @tf.function
    def fn3(self):
        self.var.assign(0)
        y = self.var
        for _ in self.dataset:
            self.var.assign(self.var+self.increment)
            y = self.var
            tf.print(y)
        tf.print(self.var)  # Inspect self.var value
        tf.print(y)
        return y

Output:

# Executed in Tensorflow 2.1.0
# test.fn3()
1
2
3
0  << self.var
3

This is already fixed If you execute this in Tensorflow 2.2.0-rc2 .
The output will be your desired outcome even when printing it during graph execution.

To quickly simulate this you could use Google Colab and use %tensorflow_version 2.x to get the latest available version for Tensorflow .

Output:

# Executed in Tensorflow 2.2.0-rc2
Function 1
1
2
3
3
Function 2
1
2
3
3
Function 3
1
2
3
3 << Value of self.var in test.fn3()
3
Function 4
1
2
3
3

You could check more about the changes in the latest Tensorflow Updates in this link .