Search code examples
pythontensorflowtensorflow2.0tensorflow-hub

TensorFlow saved_model.pb cannot be loaded correctly


I have downloaded a model from Hub. The archive contains TensorFlow 2 saved model and if decompressed shows a file named saved_model.pb and a variables folder that inside has 2 files variables.data-00000-of-00001 and variables.index.

This model seems cannot be passed to tf.keras.models.load_model() I have tried

my_model=tf.saved_model.load('extraction_path') 

but the resulting loaded_model object seems it isn't a normal model object ready to be used with something like my_model=model.predict(image)and the documentation section about the function isn't clear.

What is the proper procedure that should be used with this model format?


Solution

  • The following works for me on Python 3.9:

    import tensorflow as tf
    import numpy as np
    from PIL import Image
    
    extraction_path = "imagenet_efficientnet_v2_imagenet21k_ft1k_m_classification_2/"
    test_image_path = "test.png"
    # loading model from path
    model = tf.saved_model.load(extraction_path)
    
    # load and reshaping input image
    image_np = np.array(Image.open(test_image_path))
    input_tensor = tf.convert_to_tensor(image_np, tf.float32)
    input_tensor = input_tensor[tf.newaxis, ...]
    
    # run detection on test image
    detections = model(input_tensor)
    
    print(detections.shape)
    # TensorShape([1, 1000])
    # prediction confidence for each of the 1000 classes in the imagenet dataset