Bug: Setting `mixed_float16` raise error during model load

[ifahim]

What happened?

I am trying to load a model from Tensorflowhub using example code. It works perfect with the FP32. As soon as I add the tf.keras.mixed_precision.set_global_policy('mixed_float16') to enable mixed float, it raises an error. Looks like the dimension issue but then it works perfect with FP32

Relevant code

import tensorflow as tf
import tensorflow_hub as hub
IMAGE_SIZE = (224,224)

class_names = ['cat','dog']

#If you comment out the following line, the code works fine.
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# --------

model_handle = "https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/5"
do_fine_tuning = False
print("Building model with", model_handle)
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(model_handle, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(len(class_names),
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))
model.summary()

Relevant log output

Building model with https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/5
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [8], in <cell line: 4>()
      2 do_fine_tuning = False
      3 print("Building model with", model_handle)
----> 4 model = tf.keras.Sequential([
      5     tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
      6     hub.KerasLayer(model_handle, trainable=do_fine_tuning),
      7     tf.keras.layers.Dropout(rate=0.2),
      8     tf.keras.layers.Dense(len(class_names),
      9                           kernel_regularizer=tf.keras.regularizers.l2(0.0001))
     10 ])
     11 model.build((None,)+IMAGE_SIZE+(3,))
     12 model.summary()

File ~/miniconda3/envs/fahtx/lib/python3.8/site-packages/tensorflow/python/training/tracking/base.py:587, in no_automatic_dependency_tracking.<locals>._method_wrapper(self, *args, **kwargs)
    585 self._self_setattr_tracking = False  # pylint: disable=protected-access
    586 try:
--> 587   result = method(self, *args, **kwargs)
    588 finally:
    589   self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

File ~/miniconda3/envs/fahtx/lib/python3.8/site-packages/keras/utils/traceback_utils.py:67, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     65 except Exception as e:  # pylint: disable=broad-except
     66   filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67   raise e.with_traceback(filtered_tb) from None
     68 finally:
     69   del filtered_tb

File /tmp/__autograph_generated_fileo7avm3_o.py:74, in outer_factory.<locals>.inner_factory.<locals>.tf__call(self, inputs, training)
     72     result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)
     73 result = ag__.Undefined('result')
---> 74 ag__.if_stmt(ag__.not_(ag__.ld(self)._has_training_argument), if_body_3, else_body_3, get_state_3, set_state_3, ('result', 'training'), 1)
     76 def get_state_6():
     77     return (result,)

File /tmp/__autograph_generated_fileo7avm3_o.py:72, in outer_factory.<locals>.inner_factory.<locals>.tf__call.<locals>.else_body_3()
     70     training = False
     71 ag__.if_stmt(ag__.ld(self).trainable, if_body_2, else_body_2, get_state_2, set_state_2, ('training',), 1)
---> 72 result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)

File /tmp/__autograph_generated_fileo7avm3_o.py:72, in outer_factory.<locals>.inner_factory.<locals>.tf__call.<locals>.else_body_3.<locals>.<lambda>()
     70     training = False
     71 ag__.if_stmt(ag__.ld(self).trainable, if_body_2, else_body_2, get_state_2, set_state_2, ('training',), 1)
---> 72 result = ag__.converted_call(ag__.ld(smart_cond).smart_cond, (ag__.ld(training), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=True), fscope))), ag__.autograph_artifact((lambda : ag__.converted_call(ag__.ld(f), (), dict(training=False), fscope)))), None, fscope)

ValueError: Exception encountered when calling layer "keras_layer_3" (type KerasLayer).

in user code:

    File "/root/miniconda3/envs/fahtx/lib/python3.8/site-packages/tensorflow_hub/keras_layer.py", line 237, in call  *
        result = smart_cond.smart_cond(training,

    ValueError: Could not find matching concrete function to call loaded from the SavedModel. Got:
      Positional arguments (4 total):
        * <tf.Tensor 'inputs:0' shape=(None, 224, 224, 3) dtype=float16>
        * False
        * False
        * 0.99
      Keyword arguments: {}

     Expected these arguments to match one of the following 4 option(s):

    Option 1:
      Positional arguments (4 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='inputs')
        * True
        * True
        * TensorSpec(shape=(), dtype=tf.float32, name='batch_norm_momentum')
      Keyword arguments: {}

    Option 2:
      Positional arguments (4 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='inputs')
        * True
        * False
        * TensorSpec(shape=(), dtype=tf.float32, name='batch_norm_momentum')
      Keyword arguments: {}

    Option 3:
      Positional arguments (4 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='inputs')
        * False
        * True
        * TensorSpec(shape=(), dtype=tf.float32, name='batch_norm_momentum')
      Keyword arguments: {}

    Option 4:
      Positional arguments (4 total):
        * TensorSpec(shape=(None, None, None, 3), dtype=tf.float32, name='inputs')
        * False
        * False
        * TensorSpec(shape=(), dtype=tf.float32, name='batch_norm_momentum')
      Keyword arguments: {}

Call arguments received by layer "keras_layer_3" (type KerasLayer):
  • inputs=tf.Tensor(shape=(None, 224, 224, 3), dtype=float16)
  • training=False

​

tensorflow_hub Version

0.12.0 (latest stable release)

TensorFlow Version

other (please specify)

Other libraries

tensorflow-gpu==2.9.1

Python Version

3.x

OS

Linux

[alenarepina]

Please try passing dtype=tf.float32 to hub.KerasLayer(), i.e.

import tensorflow as tf
import tensorflow_hub as hub
IMAGE_SIZE = (224,224)

class_names = ['cat','dog']

tf.keras.mixed_precision.set_global_policy('mixed_float16')

model_handle = "https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/5"
do_fine_tuning = False
print("Building model with", model_handle)
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(model_handle, trainable=do_fine_tuning, dtype=tf.float32),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(len(class_names),
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))
model.summary()

This should fix the issue.

[sanatmpa1]

@ifahim,

Can you kindly take a look on the above comment and confirm if it resolves your question? Thank you!

[gaikwadrahul8]

@ifahim

I tried to reproduce the same error, I'm unable to get the same error after passing dtype=tf.float32 to hub.KerasLayer(), for your reference I've added Gist file here

Could you please confirm if this issue is resolved for you ? Please feel free to close the issue if it is resolved ?

Thank you!

[WGierke]

Closing due to inactivity.