Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Can't declare tf.Variable in @tf.function decorated function #26812

Closed
galeone opened this Issue Mar 17, 2019 · 9 comments

Comments

Projects
None yet
5 participants
@galeone
Copy link

commented Mar 17, 2019

System information

  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04):Archlinux
  • TensorFlow installed from (source or binary): binary
  • TensorFlow version (use command below): 2.0.0-dev20190317
  • Python version: 3.6

Describe the current behavior

A function that correctly works in eager execution can't be decorated with @tf.function if declares a tf.Variable in the function body.

The error message, reported below, is misleading since it talks about a non-first invocation when the function is invoked only once.

ValueError: tf.function-decorated function tried to create variables on non-first call.

Describe the expected behavior

Calling a function decorated with the @tf.function should produce the same output as the same function without the decoration.

Code to reproduce the issue

import tensorflow as tf

import tensorflow as tf

@tf.function
def f():
    a = tf.constant([[10, 10], [11., 1.]])
    x = tf.constant([[1., 0.], [0., 1.]])
    b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    return y

print(f())
@zakizhou

This comment has been minimized.

Copy link

commented Mar 18, 2019

see this

and rewrite your code to

import tensorflow as tf

b = None


@tf.function
def f():
    a = tf.constant([[10, 10], [11., 1.]])
    x = tf.constant([[1., 0.], [0., 1.]])
    global b
    if b is None:
        b = tf.Variable(12.)
    y = tf.matmul(a, x) + b
    return y

print(f())

or

class F(object):
    def __init__(self):
        self.b = None

    @tf.function
    def __call__(self):
        a = tf.constant([[10, 10], [11., 1.]])
        x = tf.constant([[1., 0.], [0., 1.]])
        if self.b is None:
            self.b = tf.Variable(12.)
        y = tf.matmul(a, x) + self.b
        return y


f = F()
print(f())
@galeone

This comment has been minimized.

Copy link
Author

commented Mar 18, 2019

HI @zakizhou ! I'm aware that defining a function that declares a variable inside creates a status and it has to be handled differently.

As you suggested an alternative I to wrap the function in a class and make the variable a private attribute, or declaring the variable as global and check if it is None.

However, as stated in the RFC:

State (like tf.Variable objects) are only created the first time the function f is called.

Thus, when I invoke f() for the first time, I should have the graph definition with a tf.Variable object created and executed once. And it should work according to what is written in the RFC.

Instead, what happens is that even though I call the function only once, and thus this is the first call, it seems like tf.function is evaluating the function twice (like calling f(), f()), making what is stated in the RFC false.

@alextp

This comment has been minimized.

Copy link
Member

commented Mar 19, 2019

tf.function may evaluate your python function more than once.

What the RFC states instead is that you are allowed to create variables as long as variable creation only happens the first time your python function is evaluated.

@alextp alextp closed this Mar 19, 2019

@alextp

This comment has been minimized.

Copy link
Member

commented Mar 19, 2019

And the reason for this is that if you write python code which unconditionally creates variables then I can't tell whether you mean to create a new variable every time the python function is called (eager behavior) or reuse the existing variables (what happens in graph tf 1.x) so an error felt safer.

@galeone

This comment has been minimized.

Copy link
Author

commented Mar 20, 2019

Thank you for the explanation @alextp - now everything is clearer. The RFC just describes the first call, but there is no guarantee that tf.function won't call the function more than once while converting it as a graph and for this reason, the developer should take care of handling the variables creation; thus If I need to reuse or not a variable it makes no difference, I have to take care of its status manually.

@alextp

This comment has been minimized.

Copy link
Member

commented Mar 20, 2019

@ericpts

This comment has been minimized.

Copy link

commented Mar 22, 2019

I have an exact issue like this, but with using different optimizers:

@tf.function
def apply_gradients_once(optimizer, grads, vars):
    optimizer.apply_gradients(zip(
        grads, vars))


def apply_gradients(self, use_fast, grads_per_model):
    for i in range(self.nmodels):
        grads = grads_per_model[i]
        vars = self.model[i].get_trainable_variables()

        if use_fast[i]:
            optimizer = self.fast_optimizer
            apply_gradients_once(optimizer, grads, vars)
        else:
            optimizer = self.slow_optimizer
            apply_gradients_once(optimizer, grads, vars)

On a first epoch, I have use_fast = [True, False], and then on a second one I have
use_fast = [False, True].

With this, I have a huge error trace on the apply_gradients_once call, ending with ValueError: tf.function-decorated function tried to create variables on non-first call..

From my understanding, since the function is always called with different arguments, and there are no external dependencies, this error should not be happening.

@alextp

This comment has been minimized.

Copy link
Member

commented Mar 25, 2019

@ericpts this is an interesting use case I hadn't thought of.

As a workaround I recommend you use two instances of tf.function here.

Let's open another issue to discuss this?

@alextp

This comment has been minimized.

Copy link
Member

commented Mar 25, 2019

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.