Where to apply tf.function

General Discussion
#1

At what granularity should I apply tf.function to my models and layers? Should I decorate every layer’s call with @tf.function or just the model layer’s call method, or just the training step?

It seems to me that decorating keras layers’ call method is not really intended, as tf.function does not support keyword-arguments like training or mask.

2 Likes
#2

Using tf.function gives you performance as well as reusability benefits by using same traces and gives you flexibility to use concrete functions.

Attaching doc reference for how better performance can be achieved with tf.function.

1 Like