Keras loadmodel for custom model with custom layers - Transformer documentation example

user_1506850 0
Viewed 520

I am running the following example:

https://keras.io/examples/nlp/text_classification_with_transformer/

I have created and trained a model as described, and it works nicely:

inputs = layers.Input(shape=(maxlen,))
embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
x = embedding_layer(inputs)
transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)
x = transformer_block(x,training=True)
x = layers.GlobalAveragePooling1D()(x)
x = layers.Dropout(0.1)(x)
x = layers.Dense(20, activation="relu")(x)
x = layers.Dropout(0.1)(x)
outputs = layers.Dense(2, activation="softmax")(x)

model = keras.Model(inputs=inputs, outputs=outputs)


"""
## Train and Evaluate
"""

model.compile("adam", "sparse_categorical_crossentropy", metrics=["accuracy"])
history = model.fit(
    x_train, y_train, batch_size=1024, epochs=1, validation_data=(x_val, y_val)
)

model.save('SPAM.h5')

How to save and load such custom model correctly in Keras?

I have tried

 best_model=tf.keras.models.load_model('SPAM.h5')
 ValueError: Unknown layer: TokenAndPositionEmbedding

but the model seems to miss the custom Layers. However the following also does not work

best_model=tf.keras.models.load_model('SPAM.h5',custom_objects={"TokenAndPositionEmbedding": TokenAndPositionEmbedding()})
 
TypeError: __init__() missing 3 required positional arguments:
 'maxlen', 'vocab_size', and 'embed_dim'

also passing the class does not solve.

best_model=tf.keras.models.load_model('SPAM.h5',
 custom_objects={"TokenAndPositionEmbedding": TokenAndPositionEmbedding})
 TypeError: __init__() got an unexpected keyword argument 'name'



 best_model=tf.keras.models.load_model('SPAM.h5',
{"TokenAndPositionEmbedding":
TokenAndPositionEmbedding,'TransformerBlock':TransformerBlock,
'MultiHeadSelfAttention':MultiHeadSelfAttention})
transformer-model keras tensorflow serialization python
1 Answers

Based on this answer you need to add this method (get_config) to each class (TokenAndPositionEmbedding and TransformerBlock):

TransformerBlock:

def get_config(self):
    config = super().get_config().copy()
    config.update({
        'embed_dim': self.embed_dim,
        'num_heads': self.num_heads,
        'ff_dim': self.ff_dim,
        'rate': self.rate
    })
    return config

and change the constructor to

def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
    super(TransformerBlock, self).__init__()
    self.embed_dim = embed_dim
    self.num_heads = num_heads
    self.ff_dim = ff_dim
    self.rate = rate
    self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
    self.ffn = keras.Sequential(
        [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
    )
    self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
    self.dropout1 = layers.Dropout(rate)
    self.dropout2 = layers.Dropout(rate)

TokenAndPositionEmbedding:

Similarly, add this to the class

def get_config(self):
    config = super().get_config().copy()
    config.update({
        'maxlen': self.maxlen,
        'vocab_size': self.vocab_size,
        'embed_dim': self.embed_dim
    })
    return config

and replace constructor with:

def __init__(self, maxlen, vocab_size, embed_dim, **kwargs):
    super(TokenAndPositionEmbedding, self).__init__()
    self.maxlen = maxlen
    self.vocab_size = vocab_size
    self.embed_dim = embed_dim
    self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
    self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

The reason why you cannot save and load the custom layer is explained in the answer link. When loading just do:

x = load_model('model.h5', custom_objects = {"TransformerBlock": TransformerBlock, "TokenAndPositionEmbedding": TokenAndPositionEmbedding})
user_146137680
Related