profile
viewpoint

Ask questionstf.function-decorated function tried to create variables on non-first call

System information

  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Debian Testing
  • TensorFlow installed from (source or binary): Binary
  • TensorFlow version (use command below): 2.0.0.dev20190227
  • Python version: 3.7

Describe the current behavior

A function which correctly works when in eager mode does not work anymore when annotated with tf.function.

In particular, it complains about ValueError: tf.function-decorated function tried to create variables on non-first call., even though the function is always called with different parameters.

This is a continuation of https://github.com/tensorflow/tensorflow/issues/26812#issuecomment-475600836.

Describe the expected behavior

The apply_gradients_once() function should work even when annotated with tf.function.

Code to reproduce the issue

import tensorflow as tf
import numpy as np


fast_optimizer = tf.keras.optimizers.Adam(
        learning_rate=1e-3)

slow_optimizer = tf.keras.optimizers.Adam(
        learning_rate=1e-3 * 1e-9)


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


def apply_grads(use_fast, grads_per_model, vars):
    for i in range(2):
        if use_fast[i]:
            apply_gradients_once(fast_optimizer, grads_per_model[i], vars[i])
        else:
            apply_gradients_once(slow_optimizer, grads_per_model[i], vars[i])


def compute_loss(w, x, y):
    r = (w * x - y)**2
    r = tf.math.reduce_mean(r)
    return r

def compute_gradients(model):
    with tf.GradientTape() as tape:
        tape.watch(model)
        loss = compute_loss(model, x, y)
    grads = tape.gradient(loss, model)
    return grads


w = [
    tf.Variable(0.0),
    tf.Variable(1.0)]

x = np.array([1, 2, 3])
y = np.array([1, 2, 3])

vars = []
grads = []
for i in range(2):
    vars.append([w[i]])
    grads.append(compute_gradients(w[i]))

apply_grads([True, False], grads, vars)

Other info / logs

Error log:

Traceback (most recent call last):
  File "main.py", line 52, in <module>
    apply_grads([True, False], grads, vars)
  File "main.py", line 23, in apply_grads
    apply_gradients_once(slow_optimizer, grads_per_model[i], vars[i])
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 414, in __call__
    return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 1254, in __call__
    graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 1577, in _maybe_define_function
    args, kwargs, override_flat_arg_shapes=relaxed_arg_shapes)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 1479, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 685, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 317, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 677, in wrapper
    ), args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 392, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpr2ti5o1e.py", line 4, in tf__apply_gradients_once
    ag__.converted_call('apply_gradients', optimizer, ag__.ConversionOptions(recursive=True, verbose=0, strip_decorators=(tf.function, defun, ag__.convert, ag__.do_not_convert, ag__.converted_call), force_conversion=False, optional_features=(), internal_convert_user_code=True), (zip(grads, vars),), {})
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 267, in converted_call
    return _call_unconverted(f, args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 188, in _call_unconverted
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 399, in apply_gradients
    self._create_hypers()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 558, in _create_hypers
    aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 727, in add_weight
    aggregation=aggregation)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/training/tracking/base.py", line 622, in _add_variable_with_custom_getter
    **kwargs_for_getter)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer_utils.py", line 152, in make_variable
    aggregation=aggregation)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py", line 212, in __call__
    return cls._variable_v1_call(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py", line 175, in _variable_v1_call
    aggregation=aggregation)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py", line 58, in getter
    return captured_getter(captured_previous, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 375, in invalid_creator_scope
    "tf.function-decorated function tried to create "
ValueError: tf.function-decorated function tried to create variables on non-first call.
tensorflow/tensorflow

Answer questions thisismygitrepo

The following will not work:

@tf.function
def train_one_step(model,x1,x2,target):
    optim = optimizers.Adam(0.001)  #this line means you'll try to create variables after the first call, which is not allowed
    with tf.GradientTape() as tape:
        y = model(x1,x2)
        loss = model.loss(y,target)
        grad = tape.gradient(loss,model.trainable_variables)
        clipped_gradients, _ = tf.clip_by_global_norm(grad,5)
        optim.apply_gradients(zip(grad, model.trainable_variables))
    return loss

However the following should work:

optim = optimizers.Adam(0.001)

@tf.function
def train_one_step(model,x1,x2,target):
    with tf.GradientTape() as tape:
        y = model(x1,x2)
        loss = model.loss(y,target)
    grad = tape.gradient(loss,model.trainable_variables)
    clipped_gradients, _ = tf.clip_by_global_norm(grad,5)
    optim.apply_gradients(zip(grad, model.trainable_variables))
    return loss

What you've done is basically taking the optimizer out of the train_one_step function. Suppose that I want to change the learning rate after each step, do you recommend doing it oustide, then, passing the optimizer as an argument to train_one_step ?

useful!
source:https://uonfu.com/
answerer
Alex thisismygitrepo @UQitee Brsibane, Australia. PhD. Enthusiast of math, stats, computer science.
Github User Rank List