Using mixed-precision with hub models

General Discussion
,
#1

Hi folks. When using mixed precision to perform transfer learning with any hub model I run into the following error:

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
      Positional arguments (2 total):
        * Tensor("x:0", shape=(None, 224, 224, 3), dtype=float16)
        * False
      Keyword arguments: {}
    
    Expected these arguments to match one of the following 4 option(s):
    
    Option 1:
      Positional arguments (2 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='input_1')
        * True
      Keyword arguments: {}
    
    Option 2:
      Positional arguments (2 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='x')
        * False
      Keyword arguments: {}
    
    Option 3:
      Positional arguments (2 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='input_1')
        * False
      Keyword arguments: {}
    
    Option 4:
      Positional arguments (2 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='x')
        * True
      Keyword arguments: {}

Is it a known issue? To reproduce this, just take this official example: Transfer learning with TensorFlow Hub and add the following lines of code in the library imports:

from tensorflow.keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

The data type of the last layer of the classifier should be float32. This is to prevent numeric instabilities, though. Also, to perform mixed-precision the compute capability of your GPU should be at least 7 or higher. V100 is a good example of this.

Cc: @lgusm @kempy @markdaoust

1 Like
#2

Hi Sayak,

This is expected behaviour as you cannot change number types that are saved on the SavedModel.
The model publisher could publish a version with mixed precision and allow you to do what you want but I don’t know if anyone did that yet.

From the docs: “Note: The data types used by a saved model have been fixed at saving time. Using tf.keras.mixed_precision etc. has no effect on the saved model that gets loaded by a hub.KerasLayer.”

1 Like
#3

I see. I must admit performance-wise this is indeed a huge limitation.

1 Like
#4

I don’t know if a workaround like this one could be used internally also for float32 to mixed_float16

1 Like
#5

Thanks, @Bhack. But the workaround needs us to first train the model with mixed-precision which in this case (TF-Hub) not supported it seems.

1 Like
#6

Yes I know that the example is the inverted case.
What I meant is that if TFHUB knows internally how to rebuild the model I was guessing if that workaround f32_model.set_weights(model.get_weights()) could work also in this case.

I’ve not tested this, It was just an hypotesis but if It works It could be similar to the cast in Mxnet:

You can also fine-tune a model, which was originally trained in float32, to use float16

1 Like
#7

Hey Sayak, I think if you want to see this supported in general, your best bet would be to file a GitHub issue against TensorFlow for a feature request to enable loading SavedModels with different precision than they were saved in.

2 Likes
#8

A workaround is to temporarily change the global policy to float32 when loading, then change it back to mixed precision. For example:

    tf.keras.mixed_precision.set_global_policy('float32')  # change policy

   *TRANSFER_MODEL = tf.keras.Sequential([
   *     hub.KerasLayer("https://tfhub.dev/tensorflow/efficientnet/b0/classification/1")
   *    ])

   *tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')  # change policy back
2 Likes
#9

I think it’s not a workaround, but a clear indicator that don't use mixed precision if you use tf-hub model. Very unfortunate.