SimCLR Contrastive model save gives 'EagerTensor is not JSON serializable'

[kechan]

Versions
TensorFlow: 2.8.0
TensorFlow Similarity 0.15.4

How to repro:

Ran unsupervised_hello_world.ipynb in Colab, choose 'simclr' as ALGORITHM. Execute all the cells required to instantiate a contrastive_model. Without training it, just skip to this cell

contrastive_model.save(DATA_PATH / 'models' / f'trained_model')

Errors

[/usr/local/lib/python3.7/dist-packages/tensorflow_similarity/models/contrastive_model.py](https://localhost:8080/#) in save(self, filepath, save_index, compression, overwrite, include_optimizer, save_format, signatures, options, save_traces)
    535             base_config = self.get_config()
    536             config = json.dumps(
--> 537                 {**metadata, **base_config}, cls=NumpyFloatValuesEncoder
    538             )
    539             json.dump(config, f)

[/usr/local/lib/python3.7/dist-packages/tensorflow_similarity/models/contrastive_model.py](https://localhost:8080/#) in default(self, obj)
     51         if isinstance(obj, np.float32):
     52             return float(obj)
---> 53         return json.JSONEncoder.default(self, obj)
     54 
     55

[/usr/lib/python3.7/json/encoder.py](https://localhost:8080/#) in default(self, o)
    177 
    178         """
--> 179         raise TypeError(f'Object of type {o.__class__.__name__} '
    180                         f'is not JSON serializable')
    181 

TypeError: Object of type EagerTensor is not JSON serializable

[kechan]

tf.keras.losses.serialize(contrastive_model.loss) seems to be the object that contains a tf.Tensor that's not JSON serializable:

{'class_name': 'Similarity>SimCLRLoss',
 'config': {'margin': 0.001,
  'name': 'simclr',
  'reduction': 'auto',
  'temperature': <tf.Tensor: shape=(), dtype=float32, numpy=0.5>}}

Potential fix:
In simclr.py,

def get_config(self) -> Dict[str, Any]:
    config = {
        "temperature": float(self.temperature),       #<------- cast it to float, this is json serializable.
        "margin": self.margin,
    }
    base_config = super().get_config()
    return {**base_config, **config}

Probably my TF understanding not deep enough, but there seems to be a difference in how self.temperature and self.margin are treated. self.margin is a float without any tf.constant casting like the temperature, but both are involved in the math op in this loss. Not sure if this is done for performant reason.

I will try this alteration locally, and see if there's any problem downstream.

[owenvallis]

You're right, I don't think we actually need to convert the temperature to a tensor as there isn't any performance difference. I'll change the temperature to be a normal python float. I also double checked the SimSiam and Barlow loss, and they seem to serialize properly.

[owenvallis]

Just pushed a fix in b375635. It will be pushed to pypi as 0.15.5 shortly.

[martin0258]

Hi, I got the same error in the latest version.

How to reproduce
Run unsupervised_hello_world.ipynb, with 'simsiam' as ALGORITHM.

Env versions

OS: Ubuntu 20.04.4 LTS
Python: Python 3.8.13

➜ pip list | grep 'tensor'                                                         
tensorboard                  2.9.0
tensorboard-data-server      0.6.1
tensorboard-plugin-wit       1.8.1
tensorboardX                 2.5
tensorflow                   2.9.1
tensorflow-addons            0.17.0
tensorflow-datasets          4.6.0
tensorflow-estimator         2.9.0
tensorflow-gpu               2.9.1
tensorflow-hub               0.12.0
tensorflow-io-gcs-filesystem 0.26.0
tensorflow-metadata          1.8.0
tensorflow-similarity        0.16.3

Errors

[Saving projector model]
|-path:tfsim_contrastive_model/models/trained_model/projector
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
[Saving predictor model]
|-path:tfsim_contrastive_model/models/trained_model/predictor
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
/home/martin/mambaforge/envs/tainan-search/lib/python3.8/site-packages/numpy/lib/npyio.py:518: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.
  arr = np.asanyarray(arr)
Traceback (most recent call last):
  File "unsupervised_hello_world.py", line 727, in <module>
    contrastive_model.save(DATA_PATH / "models" / f"trained_model")
  File "/home/martin/mambaforge/envs/tainan-search/lib/python3.8/site-packages/tensorflow_similarity/models/contrastive_model.py", line 534, in save
    config = json.dumps(
  File "/home/martin/mambaforge/envs/tainan-search/lib/python3.8/json/__init__.py", line 234, in dumps
    return cls(
  File "/home/martin/mambaforge/envs/tainan-search/lib/python3.8/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/home/martin/mambaforge/envs/tainan-search/lib/python3.8/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/home/martin/mambaforge/envs/tainan-search/lib/python3.8/site-packages/tensorflow_similarity/models/contrastive_model.py", line 53, in default
    return json.JSONEncoder.default(self, obj)
  File "/home/martin/mambaforge/envs/tainan-search/lib/python3.8/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type EagerTensor is not JSON serializable

[martin0258]

Potential fix:

The margin is of type tf.constant which is not serializable.

'loss': {'class_name': 'Similarity>SimSiamLoss', 
  'config': 
  {'reduction': 'auto', 
   'name': 'simsiam', 
   'projection_type': 'cosine_distance', 
   'margin': <tf.Tensor: shape=(1,), dtype=float32, numpy=array([0.001], dtype=float32)>}}

I manually edited the file simsiam.py to cast it to float, and the error is gone.

    def get_config(self) -> Dict[str, Any]:
        config = {
            "projection_type": self.projection_type,
            # "margin": self.margin,
            "margin": float(self.margin),  # Make it JSON serializable
        }
        base_config = super().get_config()
        return {**base_config, **config}

[owenvallis]

@martin0258, thanks for catching this. I'll push a patch shortly.