Training multiple models defined from the same class in Tensorflow 2.0 fails when using @tf.function

Question

I am using Tensorflow 2.1 to create custom models and custom training loops. My aim is to compare the accuracy of different configurations of my neural network. Specifically, in this case, I am comparing the reconstruction error of an AutoEncoder with varying latent dimension. Hence, I am training my network for one latent dimension then computing the test error and then I redo this process for another latent dimension, and so on. With this process I want to create plots like this:

Plot example:

Screnshot1

To speed up the training I want to use the @tf.function decorator for the BackPropagation part of my training loop. However, when I try to train several different networks, looping over the latent dimension I get an error. See below:

ValueError: in converted code:

    :41 grad  *
        loss_value = tf.losses.mean_squared_error(inputs, model(inputs))
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/engine/base_layer.py:778 __call__
        outputs = call_fn(cast_inputs, *args, **kwargs)
    :33 call  *
        x_enc = self.encoder(inp)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/engine/base_layer.py:778 __call__
        outputs = call_fn(cast_inputs, *args, **kwargs)
    :9 call  *
        x = self.dense1(inp)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/engine/base_layer.py:748 __call__
        self._maybe_build(inputs)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/engine/base_layer.py:2116 _maybe_build
        self.build(input_shapes)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/layers/core.py:1113 build
        trainable=True)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/engine/base_layer.py:446 add_weight
        caching_device=caching_device)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/training/tracking/base.py:744 _add_variable_with_custom_getter
        **kwargs_for_getter)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/engine/base_layer_utils.py:142 make_variable
        shape=variable_shape if variable_shape else None)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/ops/variables.py:258 __call__
        return cls._variable_v1_call(*args, **kwargs)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/ops/variables.py:219 _variable_v1_call
        shape=shape)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/ops/variables.py:65 getter
        return captured_getter(captured_previous, **kwargs)
    /tensorflow-2.1.0/python3.6/tensorflow_core/python/eager/def_function.py:502 invalid_creator_scope
        "tf.function-decorated function tried to create "

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

I do not get this error when I remove @tf.function decorator. I believe if it has something to do with Tensorflow creating a computational graph when I use the decorator and this graph remains when I create another instance of my network. Thus, sparking an error since the old graph does not match the new instance of the network. But I am not sure about this at all, since I believe I am missing something fundamental about Tensorflow here!

Below is a very simply version of my code recreating the error. I have tried to remove all the unnecessary parts of the code to make it easier to read and debug. Furthermore, I am generating a very simply training and test set just for the sake of this question. I have already tried the tf.keras.backend.clear_session() function without any luck.

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# Encoder
class build_encoder(tf.keras.Model):
  def __init__(self,latent_dim):
      super(build_encoder, self).__init__()

      self.dense1 = tf.keras.layers.Dense(32, activation='relu',use_bias=True)
      self.dense2 = tf.keras.layers.Dense(latent_dim, activation='relu',use_bias=True)

  def call(self, inp):
      x = self.dense1(inp)
      x = self.dense2(x)
      return x

# Decoder
class build_decoder(tf.keras.Model):
  def __init__(self,):
      super(build_decoder, self).__init__()

      self.dense1 = tf.keras.layers.Dense(32, activation='relu',use_bias=True)
      self.dense2 = tf.keras.layers.Dense(10, activation='relu',use_bias=True)

  def call(self, inp):
      x = self.dense1(inp)
      x = self.dense2(x)
      return x

# Full Autoencoder
class Autoencoder(tf.keras.Model):
  def __init__(self,latent_dim=5):
      super(Autoencoder, self).__init__()

      self.encoder = build_encoder(latent_dim)
      self.decoder = build_decoder()

  def call(self, inp):
      x_enc = self.encoder(inp)
      x_dec = self.decoder(x_enc)
      return x_dec

#### Here is the backpropagation with @tf.function decorator ####
@tf.function
def grad(model, inputs):
    with tf.GradientTape() as tape:
        loss_value = tf.losses.mean_squared_error(inputs, model(inputs))
    return loss_value, tape.gradient(loss_value, model.trainable_variables)

# Training loop function
def train(x_train, model, num_epochs, batch_size,optimizer):

    train_loss = []

    for epoch in range(num_epochs):
        tf.random.shuffle(x_train)
        for i in range(0, len(x_train), batch_size):
            x_inp = x_train[i: i + batch_size]
            loss_value, grads = grad(model, x_inp)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))

        train_loss.append(tf.reduce_mean(tf.losses.mean_squared_error(x_train, model(x_train))).numpy())

        if epoch % 100 == 0:
            print("Epoch: {}, Train loss: {:.9f}".format(epoch, train_loss[epoch]))

    return train_loss

#### Generating simple training and test data
num_train = 10000
num_test = 1000

x_train = s = np.random.uniform(0,1,(num_train,10)).astype(np.float32)
x_train[:,6:10] = 0

x_test = s = np.random.uniform(0,1,(num_test,10)).astype(np.float32)
x_test[:,6:10] = 0
###

batch_size = 8
num_epochs = 10000

test_loss = []

# Looping over the latent dimensions
for latent_dim in range(1,10):

  model = Autoencoder(latent_dim=3) # Creating an instance of my Autoencoder
  optimizer = tf.keras.optimizers.Adam(learning_rate=0.00005) # Defining an optimizer
  train_loss = train(x_train, model=model, num_epochs=num_epochs, batch_size=batch_size, optimizer=optimizer) # Training the network

  test_loss.append(tf.reduce_mean(tf.losses.mean_squared_error(x_test, model(x_test))).numpy())

plt.figure()
plt.plot(test_loss,linewidth=1.5)
plt.grid(True)
plt.show()

user_13071836 · Answer

There's an error in the code snippet you provided.

I changed last Dense layer unit from 6 to 10.

# Decoder
class build_decoder(tf.keras.Model):
  def __init__(self,):
      super(build_decoder, self).__init__()

      self.dense1 = tf.keras.layers.Dense(32, activation='relu',use_bias=True)
      self.dense2 = tf.keras.layers.Dense(10, activation='relu',use_bias=True)

  def call(self, inp):
      x = self.dense1(inp)
      x = self.dense2(x)
      return x

As for your question on training multiple model.

The error message "ValueError: tf.function-decorated function tried to create variables on non-first call" means that the function decorated by @tf.function is creating a new variable on its next iteration, this is not allowed as this function is turned into a graph.

I have modified your back propagation method, I commented out your original code to observe the difference.

#### Here is the backpropagation with @tf.function decorator ####
# @tf.function
# def grad(model, inputs):
#     with tf.GradientTape() as tape:
#         loss_value = tf.losses.mean_squared_error(inputs, model(inputs))
#     return loss_value, tape.gradient(loss_value, model.trainable_variables)

@tf.function
def MSE(y_true, y_pred):
  return tf.keras.losses.MSE(y_true, y_pred)

def backprop(inputs, model):
  with tf.GradientTape() as tape:
    loss_value = MSE(inputs, model(inputs))
  return loss_value, tape.gradient(loss_value, model.trainable_variables)

def gradient_func(model, inputs):
  return backprop(inputs, model)

The main culprit of your original code was the calling of model(inputs) as an input in the Loss Function, when you decorate @tf.function in a function it is inherited on all the functions inside, this means the Loss function is optimized.

Also a way to train multiple model without rewriting single variable, is to put them into array.

model_array = [0]
# Looping over the latent dimensions
for latent_dim in range(1,10):
  model_array.append(Autoencoder(latent_dim))
 # Creating an instance of my Autoencoder
  optimizer = tf.keras.optimizers.Adam(learning_rate=0.00005) # Defining an optimizer
  train_loss = train(x_train, model=model_array[latent_dim], num_epochs=num_epochs, batch_size=batch_size, optimizer=optimizer) # Training the network
  test_loss.append(tf.reduce_mean(tf.losses.mean_squared_error(x_test, model_array[latent_dim](x_test))).numpy())

This will arrange model into array, easier to be accessed and debugged.

Here is the complete modified code.

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# Encoder
class build_encoder(tf.keras.Model):
  def __init__(self,latent_dim):
      super(build_encoder, self).__init__()

      self.dense1 = tf.keras.layers.Dense(32, activation='relu',use_bias=True)
      self.dense2 = tf.keras.layers.Dense(latent_dim, activation='relu',use_bias=True)

  def call(self, inp):
      x = self.dense1(inp)
      x = self.dense2(x)
      return x

# Decoder
class build_decoder(tf.keras.Model):
  def __init__(self,):
      super(build_decoder, self).__init__()

      self.dense1 = tf.keras.layers.Dense(32, activation='relu',use_bias=True)
      self.dense2 = tf.keras.layers.Dense(10, activation='relu',use_bias=True)

  def call(self, inp):
      x = self.dense1(inp)
      x = self.dense2(x)
      return x

# Full Autoencoder
class Autoencoder(tf.keras.Model):
  def __init__(self,latent_dim=5):
      super(Autoencoder, self).__init__()

      self.encoder = build_encoder(latent_dim)
      self.decoder = build_decoder()

  def call(self, inp):
      x_enc = self.encoder(inp)
      x_dec = self.decoder(x_enc)
      return x_dec

#### Here is the backpropagation with @tf.function decorator ####
# @tf.function
# def grad(model, inputs):
#     with tf.GradientTape() as tape:
#         loss_value = tf.losses.mean_squared_error(inputs, model(inputs))
#     return loss_value, tape.gradient(loss_value, model.trainable_variables)

@tf.function
def MSE(y_true, y_pred):
  return tf.keras.losses.MSE(y_true, y_pred)

def backprop(inputs, model):
  with tf.GradientTape() as tape:
    loss_value = MSE(inputs, model(inputs))
  return loss_value, tape.gradient(loss_value, model.trainable_variables)

def gradient_func(model, inputs):
  return backprop(inputs, model)

# Training loop function
def train(x_train, model, num_epochs, batch_size,optimizer):

    train_loss = []

    for epoch in range(num_epochs):
        tf.random.shuffle(x_train)
        for i in range(0, len(x_train), batch_size):
            x_inp = x_train[i: i + batch_size]
            loss_value, grads = gradient_func(model, x_inp)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))
        train_loss.append(tf.reduce_mean(tf.losses.mean_squared_error(x_train, model(x_train))).numpy())

        if epoch % 100 == 0:
            print("Epoch: {}, Train loss: {:.9f}".format(epoch, train_loss[epoch]))

    return train_loss

#### Generating simple training and test data
num_train = 10000
num_test = 1000

x_train = s = np.random.uniform(0,1,(num_train,10)).astype(np.float32)
x_train[:,6:10] = 0

x_test = s = np.random.uniform(0,1,(num_test,10)).astype(np.float32)
x_test[:,6:10] = 0
###

batch_size = 8
num_epochs = 10000

test_loss = []

model_array = [0]
# Looping over the latent dimensions
for latent_dim in range(1,10):
  model_array.append(Autoencoder(latent_dim))
 # Creating an instance of my Autoencoder
  optimizer = tf.keras.optimizers.Adam(learning_rate=0.00005) # Defining an optimizer
  train_loss = train(x_train, model=model_array[latent_dim], num_epochs=num_epochs, batch_size=batch_size, optimizer=optimizer) # Training the network
  test_loss.append(tf.reduce_mean(tf.losses.mean_squared_error(x_test, model_array[latent_dim](x_test))).numpy())

plt.figure()
plt.plot(range(1,10),test_loss,linewidth=1.5)
plt.grid(True)
plt.show()

There is also a brief discussion about @tf.function and AutoGraphs in TF Documentation in this link.

Feel free to ask questions and hope this helps you.

Training multiple models defined from the same class in Tensorflow 2.0 fails when using @tf.function

1 Answers