CODEWITHSUNDEEP
pythontensorflowtensorflow2.0
william007
william007

Reputation: 18545

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

I am using tensorflow 2.3

The code below

import  tensorflow as tf

y_N= tf.Variable([1., 2., 3.],name="dd")

@tf.function
def loss():
    return -tf.reduce_mean(input_tensor=tf.reduce_sum(input_tensor=tf.math.log(y_N), axis=0))

@tf.function
def run():
    tf.keras.optimizers.Adam(0.5).minimize(loss, var_list=[y_N])

run()

gives exception

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

Problem looks like tf.keras.optimizers.Adam(0.5).minimize(loss, var_list=[y_N]) creates new variable on > first call, while using @tf.function. If must wrap it under @tf.function, how should I modify it? (in real case run() is a much bigger function)

Upvotes: 2

Views: 4359

Answers (3)

Innat
Innat

Reputation: 17219

I think it's probably not about the optimizer. Because of using function decorator, I think we need to enable the eager execution of tf.functions as follows. Source

import  tensorflow as tf
tf.config.run_functions_eagerly(True)

y_N= tf.Variable([1., 2., 3.],name="dd")

@tf.function
def loss():
    return -tf.reduce_mean(
        input_tensor=tf.reduce_sum(
            input_tensor=tf.math.log(y_N), axis=0
        )
    )

@tf.function
def run():
    tf.keras.optimizers.Adam(0.5).minimize(loss, var_list=[y_N])

run()
print(y_N)

<tf.Variable 'dd:0' shape=(3,) dtype=float32, 
numpy=array([1.5000134, 2.5000117, 3.5000103], dtype=float32)>

Update

The above code would run on eager mode with the cost of performance. To run it on graph mode, first, create instance of optimizers.Adam class and use class method afterwards as follows:

y_N = tf.Variable([1., 2., 3.],name="dd")
opt = tf.keras.optimizers.Adam(0.5)

@tf.function
def loss():
    return -tf.reduce_mean(
        input_tensor=tf.reduce_sum(
            input_tensor=tf.math.log(y_N), axis=0
        )
    )

@tf.function
def run():
    opt.minimize(loss, var_list=[y_N])

run()
print(y_N)

<tf.Variable 'dd:0' shape=(3,) dtype=float32, 
numpy=array([1.4999985, 2.499997 , 3.4999952], dtype=float32)>

Also you can check this for similar details.

Upvotes: 3

william007
william007

Reputation: 18545

I found out that after calling minimize function, it will create an iter variable, by calling self.iterations once. (This happens in tensorflow/python/keras/optimizer_v2/optimizer_v2.py)

Subsequent call of the self.iterations will not create a new variable - as iter has already been created.

Therefore, by calling to make it work, we need to create the optimizer once outside, and pass it in to @tf.function. It will still create variables in minimize function, but it will only do it once since we don't repeatedly create the Adam model.

import  tensorflow as tf

y_N= tf.Variable([1., 2., 3.],name="dd")

@tf.function
def loss():
    return -tf.reduce_mean(input_tensor=tf.reduce_sum(input_tensor=tf.math.log(y_N), axis=0))

@tf.function
def run(adam):
    adam.minimize(loss, var_list=[y_N])


adam=tf.keras.optimizers.Adam(0.5)

run(adam)

Upvotes: 0

Nicolas Gervais
Nicolas Gervais

Reputation: 36624

I'm afraid I won't be able to provide you all the explanations you might be looking for, as the documentation for @tf.function is somewhat complex, but in my experience building neural nets with Tensorflow, only the subparts of the "running" process must be decorated with @tf.function, and this part must be within with tf.GradientTape() as tape:. Your code works if I implement it that way:

import tensorflow as tf

y_N= tf.Variable([1., 2., 3.], name="dd")

@tf.function
def loss():
    return -tf.reduce_mean(input_tensor=tf.reduce_sum(
                           input_tensor=tf.math.log(y_N), axis=0))

def run():
    with tf.GradientTape() as tape:
        tf.keras.optimizers.Adam(0.5).minimize(loss, var_list=[y_N])

run()

print(y_N)

<tf.Variable 'dd:0' shape=(3,) dtype=float32, numpy=array([1.4999985, 2.499997 , 
    3.4999952], dtype=float32)>

Upvotes: 0

Related Questions