# How To Add New Keras Model During Training With TF.function

**URL:** <https://discuss.ai.google.dev/t/how-to-add-new-keras-model-during-training-with-tf-function/29505>\
**Category:** General Discussion\
**Tags:** reinforcement-learni, library, tf-function, keras\
**Created:** [February 7, 2023, 4:15am UTC](https://discuss.ai.google.dev/t/how-to-add-new-keras-model-during-training-with-tf-function/29505 "2023-02-07T04:15:06Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Avnish\_Narayan](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/avnish_narayan/32/11602_2.png) [@Avnish\_Narayan](https://discuss.ai.google.dev/u/Avnish_Narayan)\
**Post date:** [February 7, 2023, 4:15am UTC](https://discuss.ai.google.dev/t/how-to-add-new-keras-model-during-training-with-tf-function/29505/1 "2023-02-07T04:15:06Z")

</div>

Hi TF Community,

I’m a maintainer of a reinforcement learning framework called RLlib. I’ve been writing a new distributed training stack for this library using tensorflow 2.11.

One requirement that I have when training is sometimes when doing some types of Reinforcement learning training I need to add variables, or a whole new keras model, to my existing container of models that is being trained.

I’m using `tf.function` to speed up the training loop of this new distributed training stack.  
The way that I have been doing this is not by decorating my training loop function with `@tf.function` but rather holding a reference to the tf traced function by doing something like `self.update_function = tf.function(self._my_update_function)` and then calling `update_function`.

The reason that I do this is because when I add or remove keras models from my container of models, I then re-assign my update\_function in order to force retracing with the line `self.update_function = tf.function(self._my_update_function)`.

**This seems to have worked, but I’m wondering if any TF Devs or experts can chime in and let me know if there is a better way to force the retracing of my function**. I see warnings that my function has been retraced many times, after I do this retracing operation, which logically makes sense to me since I am doing a retrace operation, but then also doesn’t make sense because from a syntax standpoint it looks like I’m creating a brand new traced function object every time I do a reassignment operation and call `self.update_function = tf.function(self._my_update_function)`.

Thanks!

---

<div class="post-metadata">

**Author:** ![aniruthraj](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/aniruthraj/32/660_2.png) [@aniruthraj](https://discuss.ai.google.dev/u/aniruthraj)\
**Post date:** [October 29, 2024, 10:28am UTC](https://discuss.ai.google.dev/t/how-to-add-new-keras-model-during-training-with-tf-function/29505/3 "2024-10-29T10:28:53Z")

</div>

Hi @Avnish_Narayan,

Sorry for the delay in response.  
As far as I’m aware your approach in creating new traced function and setting this reassignment operation `self.update_function = tf.function(self._my_update_function)` is good for the RL but those warnings are due to the frequent retracing. So I recommend to reduce the model changes to minimum and reassign only when necessary and use `input_signature` to check expected input shapes, which helps in reducing unnecessary retracing.

Eg:

```auto
tf.function(
    input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])

```

Kindly review this [documentation](https://www.tensorflow.org/api_docs/python/tf/function#input_signatures) regarding `input_signature`. In addition, I recommend checking this [resource](https://www.tensorflow.org/guide/function) for more effective use of `tf.function`.

Thank You.