Expose token_embedding as a Backbone Property

[abheesht17]

• Removed self.token_embedding = token_embedding_layer.
• Rearranged all self references to follow the same order as __init__ args.

Edit:

• Exposed token_embedding as a property.

[mattdangerw]

one comment!

[mattdangerw]

These test failures are interesting, changing the layers we set on self changes our checkpoints structure. We may need to dig into this more, let's talk tomorrow!

[abheesht17]

These test failures are interesting, changing the layers we set on self changes our checkpoints structure. We may need to dig into this more, let's talk tomorrow!

Very weird. I haven't changed the names of the layers!

And it is happening only for those for which I have broken the call into two:

        # Embed tokens and positions.
        token_and_position_embedding_layer = TokenAndPositionEmbedding(
            vocabulary_size=vocabulary_size,
            sequence_length=max_sequence_length,
            embedding_dim=hidden_dim,
            embeddings_initializer=roberta_kernel_initializer(),
            name="embeddings",
        )
        embedding = token_and_position_embedding_layer(token_id_input)

instead of

        embedding = TokenAndPositionEmbedding(
            vocabulary_size=vocabulary_size,
            sequence_length=max_sequence_length,
            embedding_dim=hidden_dim,
            embeddings_initializer=roberta_kernel_initializer(),
            name="embeddings",
        )(token_id_input)

[NusretOzates]

Just a guess with a very limited log (only saw the error message or Roberta), it could be because of the self._token_embedding. Even when you don't use it, Keras tries to load its weights.

from tensorflow import keras
import tensorflow as tf


class LayerA(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.layer_b = keras.layers.Dense(1)

    def call(self, inputs):
        return self.layer_b(inputs)


class ModelA(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.layer_a = LayerA()

    def call(self, inputs):
        return self.layer_a(inputs)


a = ModelA()
a(tf.ones((1, 1)))
a.save_weights("test.h5")

after running this code, I added a new layer with self, say self.layer_b = LayerA(). But I will not use it. Even when I use the layer_b inside layer_a it will not work

from tensorflow import keras
import tensorflow as tf


class LayerA(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.layer_b = keras.layers.Dense(1)

    def call(self, inputs):
        return self.layer_b(inputs)


class ModelA(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.layer_a = LayerA()
        self.layer_b = self.layer_a.layer_b

    def call(self, inputs):
        return self.layer_a(inputs)


a = ModelA()
a(tf.ones((1, 1)))
a.load_weights("test.h5")

For now, the only solution that comes into my mind is using get_layer method of the model like this:

But with this approach, all token embeddings must have the name 'token_embedding'. I'm definitely not a fan of this approach 😄

[jbischof]

Looking great @abheesht17!

Could you update the Bert pretraining example (link) and file an Issue to update our "Getting Started" guide (link)?

[jbischof]

I agree with @NusretOzates, the self._token_embedding approach seems cleaner than relying on a specific layer name. We already have several other similar class variables.

[abheesht17]

@jbischof - it doesn't work for models with TokenAndPositionEmbedding layer. Keras considers self._token_embedding as a separate embedding layer, and errors out when we try to load preset checkpoints. Hence, this elaborate-ish solution.

[mattdangerw]

Yeah, I think what we are facing here is setattr tracking on all keras layers. Basically anytime you are setting a layer attribute on self, if gets added to a list of resources used for serialization. It looks like this can affect our checkpoint compatibility! Which is not good, we don't want to be affecting our checkpoints just to expose something like this. Relevant code -> https://github.com/keras-team/keras/blob/2727df09aa284a94ce8234ad1279d9659cdf2064/keras/engine/base_layer.py#L3215-L3229

The solution laid our here seems like a nice way to avoid the setattr tracking entirely. This LGTM.

The alternate I see would be to add a line self._auto_track_sub_layers = False to the backbone base class. But this could run us into hot water if we ever had non-functional Backbones (not everything can be a functional model -> https://keras.io/guides/functional_api/#functional-api-weakness). So the solution here seem most robust!

[mattdangerw]

I know I am dumping too much context, but for those interested in going deeper...

The __setattr__ tracking is deduped, so for Bert, where the token embedding is already a sublayer of the model directly, there is no issue here. self.some_properly = direct_layer_of_model has no issues. But Roberta for example will have the token embedding as a nested layer. self.some_property = nested_layer_of_model will change our checkpoint structure! This is what @NusretOzates was mentioning above.

Also thanks @NusretOzates for that writeup! Very helpful!

[mattdangerw]

Just the one comment re switching returning the tf.Variable to the keras.layers.Embedding

[mattdangerw]

I think we should just return the embedding layer here, that is what you are documenting on the base class below! So leave off the .embeddings part, that can be done by calling code as required.

Expose the layer seems the more general case here. It doesn't preclude getting the weights out, but also gives you a callable layer if you want it!

[mattdangerw]

Yeah, I think what we are facing here is setattr tracking on all keras layers. Basically anytime you are setting a layer attribute on self, if gets added to a list of resources used for serialization. It looks like this can affect our checkpoint compatibility! Which is not good, we don't want to be affecting our checkpoints just to expose something like this. Relevant code -> https://github.com/keras-team/keras/blob/2727df09aa284a94ce8234ad1279d9659cdf2064/keras/engine/base_layer.py#L3215-L3229

The solution laid our here seems like a nice way to avoid the setattr tracking entirely. This LGTM.

The alternate I see would be to add a line self._auto_track_sub_layers = False to the backbone base class. But this could run us into hot water if we ever had non-functional Backbones (not everything can be a functional model -> https://keras.io/guides/functional_api/#functional-api-weakness). So the solution here seem most robust!

[mattdangerw]

Thank you!

This little "cleanup PR" turned into quite a complex issue, but I def think it was good to work through!