Variable creation with different shapes in tf.function

user_3910261 0
Viewed 194

I'm trying to train a network where from one epoch to another the shapes of my trainable variables will change. So my training step looks like this

@tf.function
def training_step(inputs: tf.Tensor, 
                  labels: tf.Tensor, 
                  trainables: List[tf.Tensor], 
                  optimizer: tf.optimizers.Optimizer, 
                  loss: Callable, 
                  activation: Callable):

    def f(x, trainables):
        # do some calculation
        return result

    with tf.GradientTape() as tape:
        tape.watch(samples)
        yhat = activation(tf.vectorized_map(lambda vec: f(vec, trainables), inputs))
        loss = loss(labels,yhat)

    gradients = tape.gradient(loss, trainables)
    optimizer.apply_gradients(zip(gradients, trainables))

    return loss, trainables

When I run this for the first time, there is no problem so the first epoch runs nicely but at the beginning of the second epoch, I would like to apply some tensor decomposition which can change the shape of the trainable variables. Hence in the second epoch, I get the following error as expected

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

which is generated by this line: optimizer.apply_gradients(zip(gradients, trainables)). In this discussion it says that tf.function decorator doesn't allow running with different shapes which is understandable for efficiency reasons. But is there a way to force it to work or recompile when necessary?

PS in this link it says there is a similar bug issue but I don't think my problem is related to a bug since I'm not using regular Model module, instead, I'm simply creating a list of tf.Variable which has the tensors and then I plug it in this training_step.

Thanks

  • Python 3.6
  • Tensorflow 2.4.1
  • OS: macOS Big Sur v11.1 (also tried in Debian, same issue)
tensorflow python
1 Answers

Probably the easiest way is to use tf.function as a function instead of as a decorator. E.g.

def train_step():
    ...


some_variables = ...
train_step_compiled = tf.function(train_step)
train_step_compiled(some_variables, other_stuff)

some_variables = stuff_that_modifies_variable_shapes(some_variables)

train_step_compiled = tf.function(train_step)
train_step_compiled(some_variables, other_stuff)

I hope this "colloquial" kind of pseudo-code gets the point across -- basically you can explicitly call tf.function whenever you need it, e.g. after changing the variables, so that it will get compiled anew with the modified shapes. Calling tf.function on a function and using the returned function is the same as using it as a decorator.

user_93931020
Related