Why is predict method taking long time when used in a custom loss function?

Question

I've been trying to identify the bottleneck in my code and read about the performance issue at https://github.com/tensorflow/tensorflow/issues/40261.

My question is that when I use the predict method once the training is over, it is super fast. However, whenever I use it inside the custom loss function, it causes a serious slowdown.

import numpy as np
import random
import tensorflow as tf
from tensorflow import keras 

m = 2000
n = 250

my_input  =  tf.convert_to_tensor(np.random.random([m,n]), dtype = tf.float32) 
output =  np.random.random([m,1]) 
output = tf.convert_to_tensor(np.append(output, np.arange(output.shape[0]).reshape(-1,1), axis=1) , dtype= tf.float32) 


# some extra data to test the performance of prediction
extra_data = tf.convert_to_tensor(np.random.random([2000, n]), dtype = tf.float32) 

#use this method to make sure that tensorflow does not prune the terms that are not incorporated in the loss
@tf.function
def check(arr): 
    if arr > np.inf:
        tf.print("trivial method")

@tf.function
def obj(y_true,y_pred, model):
    loss =  tf.math.reduce_mean(tf.keras.losses.mse(tf.reshape(tf.gather(y_true, [0], axis=1), [-1,1]), y_pred)) 
    #here I use the model to make a prediction as a trivial operation
    extra = tf.math.reduce_mean(model( extra_data , training=False))
    check(extra) #If you don't use extra in a method, tensorflow removes it 

    return loss 


def my_loss(model, y_true):
    @tf.function
    def custom_loss(y_true,y_pred):
        return obj(y_true, y_pred,  model)
    return custom_loss

model = tf.keras.Sequential([  
    tf.keras.layers.Flatten(input_shape=(my_input.shape[1],)),
    tf.keras.layers.Dense(128, activation='softmax'),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(1) 
])   
             

model.compile(loss=my_loss(model, output) ,optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001))

model.fit(my_input,  output, epochs=5, batch_size=5, verbose=2)

If I comment out, extra term inside the obj method, the training does not even take a second.

Epoch 1/5
400/400 - 1s - loss: 0.1356
Epoch 2/5
400/400 - 0s - loss: 0.0832
Epoch 3/5
400/400 - 0s - loss: 0.0832
Epoch 4/5
400/400 - 0s - loss: 0.0832
Epoch 5/5
400/400 - 0s - loss: 0.0831

If I keep the term extra and use it in a trivial method (the goal here is not to impact the backpropagation to indicate that the first bottleneck is to get the predictions), then the epochs start taking way more time. Please note that this is just a toy example. In the real case, the impact is way more severe.

Epoch 1/5
400/400 - 3s - loss: 0.1365
Epoch 2/5
400/400 - 2s - loss: 0.0832
Epoch 3/5
400/400 - 2s - loss: 0.0832
Epoch 4/5
400/400 - 2s - loss: 0.0829
Epoch 5/5
400/400 - 2s - loss: 0.0827

Now, once the training is over, prediction is fairly quick.

#As dicussed in the github link, predict method is slower.
start = time.time()
check( tf.math.reduce_mean(model.predict(extra_data)))
print("Total time elapsed: ", time.time()- start) 

Total time elapsed:  0.11890172958374023


#As an alternative, using model is much faster
start = time.time()
check( tf.math.reduce_mean(model(extra_data, training=False)))
print("Total time elapsed: ", time.time()- start)   
Total time elapsed:  0.009974241256713867

I understand that this prediction operation will be repeated for every batch to complete an epoch; in other words, it will take 2000/5 = 400 repetitions, but shouldn't it be way faster since everything is done by using tf methods and objects? If I create a for loop and run model(extra_data, training=False) 400 times, then it takes total of 2 secs. I'm sure that it is not how tensorflow works under the hood. What could be the reason behind this slow execution?

Why is predict method taking long time when used in a custom loss function?

0 Answers