repeat model.call() method of tensorflow model in loop

Viewed 57

I have a tensorflow model where I need to repeatedly make 1 prediction inside a loop, and periodically I run the .fit() method. Something along the lines of

for i in range(1,n):
    x = ...
    pred = model.predict(x)
    # do something to x that depends on the value of pred
    if i % k == 0:
        model.fit(some_x_data,some_y_data,verbose=0)

The model's input layer is of shape (m,m,), so x is of size (1,m,m).

I cannot just make x a bigger numpy array/tensor with all the values I want to predict because later in each iteration of the loop I need to use that prediction to change what x will be in the next iteration of the loop when I call the predict method; so I don't know what x will be until after I run the predict on the previous x. Seems easy enough, but model.predict() has a huge amount of overhead, so it really slows things down. The fix I found for this is instead of running the predict method, use the call method instead.

model.call = tf.function(model.call,experimental_relax_shapes=True)
for i in range(n):
    x = ...
    pred = model(x,training=False)

This speeds things up by a factor of about 10 early on, but it seems as though there's a memory leak in tensorflow ( https://medium.com/dive-into-ml-ai/dealing-with-memory-leak-issue-in-keras-model-training-e703907a6501 ), such that after the loop goes on, it gets slower and slower to run the call method.

This post suggests a few fixes. Ultimately, he needed to fix something in his .fit() callback to deal with predicting on a validation set, but I don't have a validation set so I don't think this is a problem for me. The things I have tried are K.clear_session(), save-delete-reload model periodically, use a secondary model to run the call method. I have done all of these, and it seems to fix the memory leak, but when I save-delete-reload the model it seems like the weights of my adam optimizer (first and second moment of the optimizer) also get reset when I load the model. I know this because if I run

model.save('path.tf')
del model
K.clear_session()
model = tf.keras.models.load_model('path.tf')
print(model.optimizer.get_weights())

Then the output is just [].

Any suggestions of how to fix the memory leak, or get the predict method to be faster so I don't have to use the call method that creates the memory leak?

0 Answers
Related