Keras layers should override get_config to be JSON-serializable
Author: gthb
Pull request: no

# 🚀 Feature request

Support JSON serialization of Keras layers by overriding `get_config`, so that they can be sent to Tensorboard to display a conceptual graph of the model.

## Motivation

### 1. Without this, can't write model graph to Tensorboard

From https://github.com/tensorflow/tensorflow/blob/d1786ea19eb41922c0d433d71ca13b123b69b4be/tensorflow/python/ops/summary_ops_v2.py#L1004-L1009

> Writing the Keras model configuration allows the TensorBoard graph plugin to render a conceptual graph, as opposed to graph of ops. In case the model fails to serialze as JSON, it ignores and returns False.

### 2. Without this, can't save model with Keras `model.save`

The base class `get_config` method actually refuses to run if the subclass initializer has positional arguments; from `tensorflow/python/keras/engine/base_layer.py`:

```python
  @base_layer_utils.default
  def get_config(self):
    [...]
    if len(extra_args) > 1 and hasattr(self.get_config, '_is_default'):
      raise NotImplementedError('Layer %s has arguments in `__init__` and '
                                'therefore must override `get_config`.' %
                                self.__class__.__name__)
```

and all the `TF*MainLayer` classes have a `config` positional argument, so this says they “must” all override `get_config`.

And sure enough, if I make a simple Keras model using a TFBertMainLayer inside:

```python
import tensorflow as tf
from transformers import TFBertMainLayer, BertConfig


def create_model(max_sequence_len: int) -> tf.keras.Model:
    cfg = BertConfig.from_pretrained('bert-base-cased')
    bert = TFBertMainLayer(cfg)
    input_ids = tf.keras.Input(shape=(max_sequence_len,), dtype=tf.int32, name='wp_input_token_ids')
    input_mask = tf.keras.Input(shape=(max_sequence_len,), dtype=tf.bool, name='wp_input_mask')
    pooled = bert(input_ids, input_mask)[1]
    out = tf.keras.layers.Dense(units=3, activation='softmax',
                                kernel_initializer=tf.keras.initializers.glorot_uniform(),
                                use_bias=False,
                                name='classification'
                                )(pooled)
    return tf.keras.Model(inputs=[input_ids, input_mask], outputs=[out])


model = create_model(40)
model.save(filepath="tf_model.h5")
```

... then `model.save` fails:

```
Traceback (most recent call last):
  File "trysave.py", line 32, in <module>
    model.save(filepath="tf_model.h5")
  File ".../tensorflow_core/python/keras/engine/network.py", line 1008, in save
    signatures, options)
  File ".../tensorflow_core/python/keras/saving/save.py", line 112, in save_model
    model, filepath, overwrite, include_optimizer)
  File ".../tensorflow_core/python/keras/saving/hdf5_format.py", line 99, in save_model_to_hdf5
    model_metadata = saving_utils.model_metadata(model, include_optimizer)
  File ".../tensorflow_core/python/keras/saving/saving_utils.py", line 172, in model_metadata
    raise e
  File ".../tensorflow_core/python/keras/saving/saving_utils.py", line 169, in model_metadata
    model_config['config'] = model.get_config()
  File ".../tensorflow_core/python/keras/engine/network.py", line 918, in get_config
    return copy.deepcopy(get_network_config(self))
  File ".../tensorflow_core/python/keras/engine/network.py", line 1993, in get_network_config
    layer_config = serialize_layer_fn(layer)
  File ".../tensorflow_core/python/keras/utils/generic_utils.py", line 198, in serialize_keras_object
    config = instance.get_config()
  File ".../tensorflow_core/python/keras/engine/base_layer.py", line 499, in get_config
    raise NotImplementedError('Layers with arguments in `__init__` must '
NotImplementedError: Layers with arguments in `__init__` must override `get_config`.
```

## Your contribution

I got this working for the one layer I was experimenting with, like this:

```patch
diff --git a/src/transformers/modeling_tf_bert.py b/src/transformers/modeling_tf_bert.py
index 19046235..74ad621c 100644
--- a/src/transformers/modeling_tf_bert.py
+++ b/src/transformers/modeling_tf_bert.py
@@ -21,6 +21,7 @@ import logging
 import numpy as np
 import tensorflow as tf
 
+from . import PretrainedConfig
 from .configuration_bert import BertConfig
 from .file_utils import MULTIPLE_CHOICE_DUMMY_INPUTS, add_start_docstrings, add_start_docstrings_to_callable
 from .modeling_tf_utils import TFPreTrainedModel, get_initializer, shape_list
@@ -474,12 +475,20 @@ class TFBertNSPHead(tf.keras.layers.Layer):
 class TFBertMainLayer(tf.keras.layers.Layer):
     def __init__(self, config, **kwargs):
         super().__init__(**kwargs)
+        if isinstance(config, dict):
+            config = PretrainedConfig.from_dict(config)
+        self.config = config
         self.num_hidden_layers = config.num_hidden_layers
 
         self.embeddings = TFBertEmbeddings(config, name="embeddings")
         self.encoder = TFBertEncoder(config, name="encoder")
         self.pooler = TFBertPooler(config, name="pooler")
 
+    def get_config(self):
+        cfg = super().get_config()
+        cfg['config'] = self.config.to_dict()
+        return cfg
+
     def get_input_embeddings(self):
         return self.embeddings

```

and I didn't need to modify any other layer classes, just the main layer.

So maybe it's enough to do this for all the `MainLayer` classes:

```
$ rg 'class .*MainLayer\(tf.keras.layers.Layer\)' src | cat
src/transformers/modeling_tf_openai.py:class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_transfo_xl.py:class TFTransfoXLMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_xlm.py:class TFXLMMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_xlnet.py:class TFXLNetMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_distilbert.py:class TFDistilBertMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_bert.py:class TFBertMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_albert.py:class TFAlbertMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_ctrl.py:class TFCTRLMainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_t5.py:class TFT5MainLayer(tf.keras.layers.Layer):
src/transformers/modeling_tf_gpt2.py:class TFGPT2MainLayer(tf.keras.layers.Layer):
```

... or, neater, to extract a single `TFMainLayer(tf.keras.layers.Layer)` superclass for all of them, to do this in one place.

Comments:
Thanks a lot for investigating this and submitting a fix, it's awesome.

I respond in the PR it-self