huggingface optimizer cannot de-serialize
Author: suhasvk
Pull request: no
Labels: name: wontfix

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): Bert

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

    >>> import transformers
    >>> opt, _ = transformers.optimization_tf.create_optimizer(init_lr=2e-5, num_train_steps=1000, num_warmup_steps=10)
    >>> opt.__class__.from_config(opt.get_config())
    Traceback (most recent call last):
        File "<stdin>", line 1, in <module>
        File "/home/ec2-user/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 741, in from_config
        config["learning_rate"], custom_objects=custom_objects)
        File "/home/ec2-user/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/keras/optimizer_v2/learning_rate_schedule.py", line 992, in deserialize
        printable_module_name="decay")
        File "/home/ec2-user/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 292, in deserialize_keras_object
        config, module_objects, custom_objects, printable_module_name)
        File "/home/ec2-user/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/keras/utils/generic_utils.py", line 250, in class_and_config_for_serialized_keras_object
        raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    ValueError: Unknown decay: WarmUp

## Expected behavior

One would expect that the `tf.keras.optimizers.Adam` object `opt` is successfully reconstructed from its config dictionary.

After inspecting `transformers.optimization_tf`, what I believe is happening here is that `create_optimizer` returns a regular `tf.keras.optimizers.Adam` object in the event that the `weight_decay_rate` parameter is zero:

    91 else:
    92    optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule, epsilon=adam_epsilon)

However, if `num_warmup_steps` is nonzero, this object is instantiated with a _custom_ `tf.keras.optimizers.schedules.LearningRateSchedule` object of type `WarmUp`:

    77 if num_warmup_steps:
    78     lr_schedule = WarmUp(
    79         initial_learning_rate=init_lr, decay_schedule_fn=lr_schedule, warmup_steps=num_warmup_steps,
    80     )

In this case the `tf.keras.optimizers.Adam` implementation of `from_config` does not know to pass this class under the `custom_objects` keyword argument, causing deserialization to fail.

In my specific setting, this causes the method `horovod.create_distributed_optimizer(opt)` to fail, since it relies on serialization / deserialization to programmatically extend the passed optimizer class.

A solution that should work in my specific setting (and more generally) is to create an intermediate class such as

        class AdamWarmUp(tf.keras.optimizers.Adam):
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)  
            @classmethod
            def from_config(cls, config):
                custom_objects = {"WarmUp":WarmUp}
                return super().from_config(config, custom_objects=custom_objects)

which is returned by `create_optimizer` and correctly implements `tf.keras.optimizers.Optimizer`, and to have the `AdamWeightDecay` optimizer extend this class. 

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.0.0
- Platform: Linux-4.14.181-108.257.amzn1.x86_64-x86_64-with-glibc2.9
- Python version: 3.6.10
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): 2.1.0 (True)
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: True


Comments:
Great, thanks for your analysis! Do you think you can open a PR with your proposed fix?

Sure, will do!

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Hi @suhasvk , Did you open a PR for the fix. I encountered the same issue. 

Hi @LysandreJik , is there any PR for the fix. I encountered the same issue.