Lookahead bug when loading saved model

[GaranceRichard]

System information

• OS Platform and Distribution: Windows 10
• TensorFlow version and how it was installed (source or binary): TF 2.1 from pip
• TensorFlow-Addons version and how it was installed (source or binary): 0.8.3
• Python version: python 3.7
• Is GPU used? (yes/no): yes cuda 10.1

Describe the bug

My model uses TFA: Lookahead with Radam (ranger)
Trained it: ok
Saved it: ok
Impossible to load it:

Traceback (most recent call last):
  File "Boucle_IA.py", line 24, in <module>
    model = tf.keras.models.load_model('inference_complete-2020-03-18-Mobilenet.h5', custom_objects={'Lookahead': ranger})
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\saving\save.py", line 146, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\saving\hdf5_format.py", line 184, in load_model_from_hdf5
    training_config, custom_objects))
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\saving\saving_utils.py", line 229, in compile_args_from_training_config
    optimizer_config, custom_objects=custom_objects)
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\optimizers.py", line 819, in deserialize
    printable_module_name='optimizer')
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\utils\generic_utils.py", line 303, in deserialize_keras_object
    list(custom_objects.items())))
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_addons\optimizers\lookahead.py", line 185, in from_config
    custom_objects=custom_objects,
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\optimizers.py", line 819, in deserialize
    printable_module_name='optimizer')
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\utils\generic_utils.py", line 303, in deserialize_keras_object
    list(custom_objects.items())))
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\optimizer_v2\optimizer_v2.py", line 733, in from_config
    return cls(**config)
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 809, in wrapper
    check_argument_types(memo)
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 670, in check_argument_types
    raise exc from None
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 668, in check_argument_types
    check_type(description, value, expected_type, memo)
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 598, in check_type
    format(argname, qualified_name(expected_type), qualified_name(value)))
TypeError: type of argument "total_steps" must be int; got float instead

Code to reproduce the issue

on definition :

radam = tfa.optimizers.RectifiedAdam(
    lr=1e-3,
    total_steps=10000,
    warmup_proportion=0.1,
    min_lr=1e-5,
)
ranger = tfa.optimizers.Lookahead(radam, sync_period=6, slow_step_size=0.5)

on load :

import numpy as np 
import tensorflow as tf 
import tensorflow_addons 
import os
import cv2 

from pathlib import Path

# Constantes

print('loading model...')
radam = tfa.optimizers.RectifiedAdam(lr=1e-3, total_steps=10000, warmup_proportion=0.1, min_lr=1e-5)

ranger = tfa.optimizers.Lookahead(radam, sync_period=6, slow_step_size=0.5, name='Lookahead',)

model = tf.keras.models.load_model('inference_complete-2020-03-18-Mobilenet.h5', custom_objects={'Lookahead': ranger})
print('model loaded')

Tips: When I tried to open H5 copy of the saved model: in total_steps, value = 10000.0

[Squadrick]

cc @CyberZHG

[gabrieldemarmiesse]

I opened a pull request to fix the bug. The fix will be backported to the 0.8 branch. @GaranceRichard when this is merged, you'll be able to load your model using tfa-nightly. When the backport is merged and a release is done, you'll be able to load your model with tensorflow-addons 0.8.4.

[GaranceRichard]

found a solution to load it :

load without compilation, import optimizers, compile with optimizers...

it works but it could be better...

[RobeeF]

Still have the problem after updating the package with the last version...

[Lif3line]

This bug appears to still exist in 0.13 at least when saving/loading a model via MLFlow.

[bhack]

If you have a minimal code snippet to reproduce this with 0.13 please open a new ticket

[Lif3line]

The issue disappears in a minimal example I mocked up. The problem likely comes from the distributed setup I'm using, I'm guessing from inadvertently saving in 0.12 followed by extracting in 0.13. I'll raise separately if it's an issue that can be reproduced using only 0.13.