Keras V3 Automatic Speech Recognition with Transformer を試す。

2. 学習を途中で中断できるように、weight.h5 を保存できるように、手を加える。
これが、従来の方法でやると、なかなかできなかった。

結論を先に述べると、
keras.layers.Layer を継承した Custom Layer を組み込む場合、
Making new layers and models via subclassing の、
"You can optionally enable serialization on your layers" の欄の様に、
class Linear(keras.layers.Layer): .... def build(self, input_shape): .... def get_config(self): ....
build() と get_config() の2メソッドを自分で追加しないといけないみたいぞね。
従来は、 def __init__() ですべて済ませたいた部分を、
def __init__() で、保存したいクラス変数を、self.xxx=xxx ですべて保存して、
実際の割当を、def build() で行う様に分けるみたい。

def get_config() は、model.keras や weight.h5 を保存するときに、保存するクラス変数を返す。

class TokenEmbedding(layers.Layer): の変更例。
class TokenEmbedding(layers.Layer): def __init__(self, num_vocab=1000, maxlen=100, num_hid=64): super().__init__() self.num_vocab=num_vocab self.maxlen = maxlen self.num_hid=num_hid #self.emb = tf.keras.layers.Embedding(num_vocab, num_hid) #self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=num_hid) # add by nishi def build(self, input_shape): self.emb = tf.keras.layers.Embedding(self.num_vocab, self.num_hid) self.pos_emb = layers.Embedding(input_dim=self.maxlen, output_dim=self.num_hid) def call(self, x): maxlen = tf.shape(x)[-1] x = self.emb(x) positions = tf.range(start=0, limit=maxlen, delta=1) positions = self.pos_emb(positions) return x + positions ### Start 追加されたコード def get_config(self): config = { "num_vocab": self.num_vocab, "maxlen": self.maxlen, "num_hid": self.num_hid, } base_config = super().get_config() return dict(list(base_config.items()) + list(config.items())) ### End

やれやれぞね!!

参照。
transformers / tasks / asr