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

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
  • TensorFlow installed from (source or binary): pip
  • TensorFlow version (use command below): 2.0.0
  • Python version: 3.6.7
  • CUDA/cuDNN version: 7.6.0
  • GPU model and memory: 1660Ti

Describe the current behavior Error raised when I tried to train a model with function decorated with @tf.function by two different optimizer sequentially ValueError: tf.function-decorated function tried to create variables on non-first call

Describe the expected behavior The model can be trained by two different optimizer sequentially

Code to reproduce the issue Provide a reproducible test case that is the bare minimum necessary to generate the problem.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Flatten, Dense
from tensorflow.keras import Model

try: 
    gpus= tf.config.experimental.list_physical_devices('GPU')
    tf.config.experimental.set_memory_growth(gpus[0], True)
except:
    print('Failed to configure GPU')

cifar10 = tf.keras.datasets.cifar10

(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train = x_train / 255.0


class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = Flatten()
        self.d1 = Dense(20, activation='relu')
        self.d2 = Dense(10, activation='softmax')

    def call(self, x):
        x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)


optimizer = keras.optimizers.Adam(0.001)
optimizer2 = keras.optimizers.SGD(0.01)

model = MyModel()


@tf.function  # it will work correctly if this line is disabled but the training becomes very slow
def train_step(model, data, labels, optimizer):
    with tf.GradientTape() as tape:
        pred = model(data)
        loss = keras.losses.SparseCategoricalCrossentropy()(labels, pred)

    grads = tape.gradient(loss, model.trainable_variables)

    optimizer.apply_gradients(zip(grads, model.trainable_variables))


train_step(model, x_train[:10], y_train[:10], optimizer)  # this works
train_step(model, x_train[:10], y_train[:10], optimizer2)  # it can't work

Other info / logs

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-7-f4e941ca2116> in <module>
      1 train_step(model, x_train[:10], y_train[:10], optimizer)
----> 2 train_step(model, x_train[:10], y_train[:10], optimizer2)

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py in __call__(self, *args, **kwds)
    455 
    456     tracing_count = self._get_tracing_count()
--> 457     result = self._call(*args, **kwds)
    458     if tracing_count == self._get_tracing_count():
    459       self._call_counter.called_without_tracing()

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py in _call(self, *args, **kwds)
    485       # In this case we have created variables on the first call, so we run the
    486       # defunned version which is guaranteed to never create variables.
--> 487       return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
    488     elif self._stateful_fn is not None:
    489       # Release the lock early so that multiple threads can perform the call

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in __call__(self, *args, **kwargs)
   1820   def __call__(self, *args, **kwargs):
   1821     """Calls a graph function specialized to the inputs."""
-> 1822     graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
   1823     return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
   1824 

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   2148         graph_function = self._function_cache.primary.get(cache_key, None)
   2149         if graph_function is None:
-> 2150           graph_function = self._create_graph_function(args, kwargs)
   2151           self._function_cache.primary[cache_key] = graph_function
   2152         return graph_function, args, kwargs

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   2039             arg_names=arg_names,
   2040             override_flat_arg_shapes=override_flat_arg_shapes,
-> 2041             capture_by_value=self._capture_by_value),
   2042         self._function_attributes,
   2043         # Tell the ConcreteFunction to clean up its graph once it goes out of

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    913                                           converted_func)
    914 
--> 915       func_outputs = python_func(*func_args, **func_kwargs)
    916 
    917       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    356         # __wrapped__ allows AutoGraph to swap in a converted function. We give
    357         # the function a weak reference to itself to avoid a reference cycle.
--> 358         return weak_wrapped_fn().__wrapped__(*args, **kwds)
    359     weak_wrapped_fn = weakref.ref(wrapped_fn)
    360 

~/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py in wrapper(*args, **kwargs)
    903           except Exception as e:  # pylint:disable=broad-except
    904             if hasattr(e, "ag_error_metadata"):
--> 905               raise e.ag_error_metadata.to_exception(e)
    906             else:
    907               raise

ValueError: in converted code:

    <ipython-input-2-db5119c2db0a>:27 train_step  *
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:433 apply_gradients
        _ = self.iterations
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:541 __getattribute__
        return super(OptimizerV2, self).__getattribute__(name)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:655 iterations
        aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:805 add_weight
        aggregation=aggregation)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/training/tracking/base.py:744 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer_utils.py:139 make_variable
        shape=variable_shape if variable_shape else None)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/ops/variables.py:258 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/ops/variables.py:219 _variable_v1_call
        shape=shape)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/ops/variables.py:65 getter
        return captured_getter(captured_previous, **kwargs)
    /home/wilson/venv_tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py:413 invalid_creator_scope
        "tf.function-decorated function tried to create "

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

How should I correct it if possible? Or it’s still a bug? related issue #27120 Thanks!

About this issue

  • Original URL
  • State: closed
  • Created 5 years ago
  • Comments: 17 (1 by maintainers)

Most upvoted comments

I fix this problem by adding tf.config.experimental_run_functions_eagerly(True) after import tensorflow.

the following links may help:

+33
Papageno2 on Mar 8, 2020

Issue still exists in latest 2.3.1.

Running function eagerly is not a solution.

+14
pc2005 on Sep 25, 2020

I fix this problem by adding tf.config.experimental_run_functions_eagerly(True) after import tensorflow.

the following links may help:

This is equivalent to remove the tf.function decorator, this is unwanted, cf. hgffly’s comment: @tf.function # it will work correctly if this line is disabled but the training becomes very slow

+10
nicolaspi on May 26, 2020

I fix this problem by adding tf.config.experimental_run_functions_eagerly(True) after import tensorflow.

the following links may help:

This is great, it worked! Thanks.

+4
georgeliu1998 on Sep 11, 2020

I have tried the latest version TF2.1.0rc today, but it still raises errors. Maybe we need to wait longer… XD

+4
hgffly on Dec 10, 2019

If you don’t want to run functions eagerly, add a wrapper, or remove the tf.function decorator, you may find the following a better suggestion:

optimizer = tf.keras.optimizers.Adam(0.001)
optimizer2 = tf.keras.optimizers.SGD(0.01)

@tf.function 
def train_step_with_opt_1(model, data, labels):
    with tf.GradientTape() as tape:
        pred = model(data)
        loss = tf.keras.losses.SparseCategoricalCrossentropy()(labels, pred)

    grads = tape.gradient(loss, model.trainable_variables)

    optimizer.apply_gradients(zip(grads, model.trainable_variables))

@tf.function 
def train_step_with_opt_2(model, data, labels):
    with tf.GradientTape() as tape:
        pred = model(data)
        loss = tf.keras.losses.SparseCategoricalCrossentropy()(labels, pred)

    grads = tape.gradient(loss, model.trainable_variables)

    optimizer2.apply_gradients(zip(grads, model.trainable_variables))

train_step_with_opt_1(model, x_train[:10], y_train[:10]) 
train_step_with_opt_2(model, x_train[:10], y_train[:10])

The optimizer creates variables when apply_gradients is called for the first time on it. tf.keras.optimizers.SGD which is a “stateless” optimizer still creates variables (you can check by calling optimizer2.variables() before and after training, and you will see that there is a iter tf.Variable). Thus, you should create two tf.functions, one for each optimizer, if you would like to train your model with both optimizers.

I hope this is helpful!

+3
monicadsong on Dec 12, 2020

I fix this problem by adding tf.config.experimental_run_functions_eagerly(True) after import tensorflow.

the following links may help:

Thank you. It’s works well!

+3
Github-XKou on Sep 25, 2020

I fix this problem by adding tf.config.experimental_run_functions_eagerly(True) after import tensorflow. the following links may help:

This is equivalent to remove the tf.function decorator, this is unwanted, cf. hgffly’s comment: @tf.function # it will work correctly if this line is disabled but the training becomes very slow

eagerly = True equals do not use @tf.function

+2
linmeon on Jul 2, 2020

I fix this problem by adding tf.config.experimental_run_functions_eagerly(True) after import tensorflow.

the following links may help:

Tested v2.1 and fixes the issue

+1
Riathoir on Mar 13, 2020
Read more comments on GitHub
← tensorflow: ValueError: No gradients provided for any variable (Keras 2.4, Tensorflow 2.3.0)
tensorflow: ValueError: Unable to save the object ListWrapper(...) (a list wrapper constructed to track trackable TensorFlow objects) when calling the method tf.keras.Model.save_weights →