tensorflow function in graph mode (tf.function) is slow when return object

user_10545499 0
Viewed 151

I have a function wrapped by @tf.function, the function was quick when it doen't return any object. But it becomes significant slower when it return an object.

Any way to boost the performance even when it has to return an object?

Below is a toy example

from tqdm import tqdm
import tensorflow as tf
import numpy as np

model = tf.keras.Sequential([
    tf.keras.layers.Input((4,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10000)
])

model.compile(
        optimizer="rmsprop",
        loss="sparse_categorical_crossentropy",
        metrics=["sparse_categorical_accuracy"],
)

np.random.seed(0)
xy = tf.data.Dataset.from_tensor_slices((np.random.random((200000,4)), np.random.randint(1000, size=(200000,1)))).batch(200)

@tf.function
def _no_return(batch):
    x, y = batch
    y_pred = model(x, training=False)
    r = tf.math.top_k(y_pred, k=20).indices
    # True indicate that item is the correct prediction
    r = tf.cast(tf.equal(r, tf.expand_dims(tf.cast(y, tf.int32), 1)), tf.float32)
    # rank of the correct prediction, rank = 9999999+1 if no correction prediction within topk
    r = tf.add((tf.reduce_sum(r, 1)-1) * -9999999, tf.cast(tf.argmax(r, 1) + 1, tf.float32))

@tf.function
def _return(batch):
    x, y = batch
    y_pred = model(x, training=False)
    r = tf.math.top_k(y_pred, k=20).indices
    # True indicate that item is the correct prediction
    r = tf.cast(tf.equal(r, tf.expand_dims(tf.cast(y, tf.int32), 1)), tf.float32)
    # rank of the correct prediction, rank = 9999999+1 if no correction prediction within topk
    r = tf.add((tf.reduce_sum(r, 1)-1) * -9999999, tf.cast(tf.argmax(r, 1) + 1, tf.float32))
    return r

#
# this run 1428.27 iterations/s 
#
for batch in tqdm(xy):
    _no_return(batch)

#
# this run 135.61 iterations/s 
#
for batch in tqdm(xy):
    _return(batch)
tensorflow numpy python
1 Answers

It is happening because of autograph. In the case of the _return function, because the r value is returned, it is generating the complete graph (with r as the leaf node). But, in the case of _no_return function, there is no return value, and the autograph does not set r as the leaf node, and because there is no leaf node the overall graph has nothing in it.

Thus, you observe this speed-up because the function is not computing anything.

Add any leaf nodes (as shown in the demo below)

def _no_return(batch):
    # your old code
    tf.print(r) # this will create another leaf node, 
    # or you can try
    tf.summary.scalar(r)

you will observe a similar performance between _return and _no_return. Also, if you remove @tf.function decorator, then also you will see the same performance.

lemme know if you have any more doubts.

user_42844070
Related