使用字典和 for 循环调用 tf2 模型时出现意外警告

Question

In Tensorflow V2, when I call tf models inside a dictionary, it yields a warning. This is the dictionary of models:

 import tensorflow as tf
 import numpy as np

width = 128
height = 128
nchannels = 10
nclasses = 5

models = {   
    'MobileNetV2': tf.keras.applications.MobileNetV2,                                                      
    'DenseNet121': tf.keras.applications.DenseNet121,
    'DenseNet169': tf.keras.applications.DenseNet169,
    'DenseNet201': tf.keras.applications.DenseNet201,
    'InceptionResNetV2': tf.keras.applications.InceptionResNetV2,
    'InceptionV3': tf.keras.applications.InceptionV3,
    'MobileNet': tf.keras.applications.MobileNet,
    'Xception': tf.keras.applications.Xception,
    'NASNetLarge': tf.keras.applications.NASNetLarge
    }


for name, arch in models.items():

   inp = np.random.rand(1, width, height, nchannels)
   print(inp.shape) 

   model_arch = arch(input_shape=[width, height, nchannels],
                                                   include_top=False,
                                                   weights=None)
  
   model_classification = tf.keras.layers.Dense(nclasses)

   model = tf.keras.Sequential([
    model_arch,
    tf.keras.layers.Flatten(),
    model_classification])

   res = model.predict(inp)
   print(name, res.shape)

And this is the warning:

WARNING:tensorflow:7 out of the last 7 calls to <function 
Model.make_predict_function.<locals>.predict_function at 0x0000019FA812C790> 
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/tutorials/customization
/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs
/python/tf/function for  more details.

How can I avoid the warning? What is the proper way to train multiple tf models on the same dataset?

Answer 1

The issue comes from multiple model initializations at once inside the models dictionary definition. Since you are going to train one model at a time, it is better to defer the model initialization to the inner loop over the models. Adding one dense layer for classification, your model definition could look like this:

N_CLASSES = 48  # Define your own, based on your dataset!

models = {
    'VGG16': tf.keras.applications.VGG16,
             ....
    'VGG19': tf.keras.applications.VGG19  
    }

for name, arch in models.items():
   # You may have to define a custom input size, according yo tour dataset.
   #  Use  arch(input_size=(WIDTH, HEIGHT, N_CHANNELS)) for that.
   model_arch = arch()
   model_classification = tf.keras.layers.Dense(N_CLASSES)
   
   model = tf.keras.Sequential([
        model_arch,
        tf.keras.layers.Flatten(),
        model_classification
   ])
        
   # Training loop, including the data pipeline
   # Stop training, measure accuracy and store model.
   #   -> Next iteration will erase whatever is not stored.

IMPORTANT! You have to go through the training loop (eg: model.fit(train_data) ) in order to avoid the warning. Otherwise Tensorflow will still warn you that you're being very inefficient when generating an entire model just to run a single prediction (what your code does).