JSONDecodeError When Loading Model on Kaggle TPU TF 2.4.1

[MarkWijkhuizen]

When loading the model in a Kaggle TPU notebook with Tensorflow version 2.4.1 a JSONDecodeError is thrown.

What would be the cause of this error and how could one load a ConvNext model in Tensorflow 2.4.1?

print(f'tensorflow version: {tf.__version__}')
print(f'tensorflow keras version: {tf.keras.__version__}')
print(f'python version: P{sys.version}')
​
model_gcs_path = "gs://tfhub-modules/sayakpaul/convnext_tiny_1k_224/1/uncompressed"
model = tf.keras.models.load_model(model_gcs_path)
print(model.summary(expand_nested=True))

Gives the following error trace:

tensorflow version: 2.4.1
tensorflow keras version: 2.4.0
python version: P3.7.10 | packaged by conda-forge | (default, Feb 19 2021, 16:07:37) 
[GCC 9.3.0]

JSONDecodeError                           Traceback (most recent call last)
/tmp/ipykernel_48/3485362230.py in <module>
      4 
      5 model_gcs_path = "gs://tfhub-modules/sayakpaul/convnext_tiny_1k_224/1/uncompressed"
----> 6 model = tf.keras.models.load_model(model_gcs_path)
      7 print(model.summary(expand_nested=True))

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py in load_model(filepath, custom_objects, compile, options)
    210       if isinstance(filepath, six.string_types):
    211         loader_impl.parse_saved_model(filepath)
--> 212         return saved_model_load.load(filepath, compile, options)
    213 
    214   raise IOError(

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in load(path, compile, options)
    136   # Recreate layers and metrics using the info stored in the metadata.
    137   keras_loader = KerasObjectLoader(metadata, object_graph_def)
--> 138   keras_loader.load_layers(compile=compile)
    139 
    140   # Generate a dictionary of all loaded nodes.

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in load_layers(self, compile)
    374       self.loaded_nodes[node_metadata.node_id] = self._load_layer(
    375           node_metadata.node_id, node_metadata.identifier,
--> 376           node_metadata.metadata)
    377 
    378     for node_metadata in metric_list:

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _load_layer(self, node_id, identifier, metadata)
    395   def _load_layer(self, node_id, identifier, metadata):
    396     """Load a single layer from a SavedUserObject proto."""
--> 397     metadata = json_utils.decode(metadata)
    398 
    399     # If node was already created

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/json_utils.py in decode(json_string)
     67 
     68 def decode(json_string):
---> 69   return json.loads(json_string, object_hook=_decode_helper)
     70 
     71 

/opt/conda/lib/python3.7/json/__init__.py in loads(s, encoding, cls, object_hook, parse_float, parse_int, parse_constant, object_pairs_hook, **kw)
    359     if parse_constant is not None:
    360         kw['parse_constant'] = parse_constant
--> 361     return cls(**kw).decode(s)

/opt/conda/lib/python3.7/json/decoder.py in decode(self, s, _w)
    335 
    336         """
--> 337         obj, end = self.raw_decode(s, idx=_w(s, 0).end())
    338         end = _w(s, end).end()
    339         if end != len(s):

/opt/conda/lib/python3.7/json/decoder.py in raw_decode(self, s, idx)
    353             obj, end = self.scan_once(s, idx)
    354         except StopIteration as err:
--> 355             raise JSONDecodeError("Expecting value", s, err.value) from None
    356         return obj, end

JSONDecodeError: Expecting value: line 1 column 1 (char 0)

[Sayak Paul]

Does it load without TPUs?

[MarkWijkhuizen]

This error occurs generally in Tensorflow 2.4 (2.4.1 specifically), with a GPU/TPU/CPU.

[Sayak Paul]

Have you tried providing custom objects during loading as described here?

https://www.tensorflow.org/guide/keras/save_and_serialize#custom_layer_and_function_example

[MarkWijkhuizen]

Hi @sayakpaul

Thanks for your quick response.
I have difficulty understanding how to correctly load the model with the link you provided.
How should the original code be changed to load the model with the method you referenced?

model_gcs_path = "gs://tfhub-modules/sayakpaul/convnext_tiny_1k_224/1/uncompressed"
model = tf.keras.models.load_model(model_gcs_path)
print(model.summary(expand_nested=True))

[Sayak Paul]

Sorry for not being clear earlier.

If you take a look at https://github.com/sayakpaul/ConvNeXt-TF/blob/main/models/convnext_tf.py, you'll notice that there are a couple of layers (StochasticDepth for example) that are custom implemented. My guess is that the specific TF version that you are using is unable to recognize those custom layers from the SavedModel.

In order to get around the problem, we need to define a dictionary mapping the custom layer names to their original objects as shown in the link above.

Does this make sense?

[Sayak Paul]

Closing this issue due to inactivity. Feel free to re-open.