python - Is there a way to batch predict with multiple models?

Question

I'm using TensorFlow for a reinforcement learning project in which I'm creating a population of models that predict the output for a large dataset of inputs. After each model is evaluated, the scores are calculated and the models with the lowest score are removed from the population and replaced with tweaked-parameter versions of the best networks.

I am able to use batch prediction by running the dataset on individual members of the population, but I am unable to avoid the loop in which I run model.predict on every member of the population:

for model in population:
    outputs = model.predict(inputs, batch_size=BATCH_SIZE)

Since model.predict is a tf.function, I get the following warning every time I run it:

WARNING:tensorflow:5 out of the last 15 calls to <function Model.make_predict_function
<locals>.predict_function at 0x7f99d5e41730> triggered tf.function retracing. Tracing is
expensive and the excessive number of tracings could be due to (1) creating @tf.function
repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects
instead of tensors. For (1), please define your @tf.function outside of the loop. For (2),
@tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can
avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guid
/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for 
more details.

I'm not that experienced from TensorFlow, but from what I understand, a tf.function relies on tracing functions into graphs in order to make the runtime much quicker. Furthermore, by running model.predict on each model so many times, I am using a lot of compute power in retracing each model.

Is there a way to make use of batch training on multiple models? I definitely feel like I'm doing something wrong if I'm getting spammed this warning every time I call the function. For now I have disabled TensorFlow warning logging, but I'm still uneasy about this situation.

question from:https://stackoverflow.com/questions/66052522/is-there-a-way-to-batch-predict-with-multiple-models

Answer 1

Waitting for answers