TensorFlowメモ(CudnnLSTM)
DeepSpeechでは、「tensorflow.contrib」モジュールの「CudnnLSTM」が使われているので調べてみる。
「CudnnLSTM」のソースコード
tensorflow/cudnn_rnn.py at v1.15.2 · tensorflow/tensorflow · GitHub
実行する環境は「TensorFlow 2.X」系
% python -c 'import tensorflow as tf; print(tf.__version__)' 2.1.0
「LSTM」を使った実装
まずは比較用に、「LSTM」を使った実装での動きを確認する。4つの数字の羅列をもとに、クラス分けを行うテスト用プログラムを用意。
- インプットが「1, 2, 3, 4」だったら、クラス「0」
- インプットが「3, 4, 5, 6」だったら、クラス「1」
- インプットが「5, 6, 7, 8」だったら、クラス「2」
pythonコード
import tensorflow as tf import numpy as np # Input Data train_images = [[1, 2, 3, 4], [3, 4, 5, 6], [5, 6, 7, 8]] train_images = tf.convert_to_tensor(train_images, np.int32) num_class = 3 train_labels = [0, 1, 2] categorical_labels = tf.keras.utils.to_categorical(train_labels, num_classes=num_class)
「Embedding」、「LSTM」、「Dense」で構成される3層モデル。
pythonコード(続き)
# Define Model def create_model(vocab_size=10): model = tf.keras.Sequential([ tf.keras.layers.Embedding(vocab_size, 80), tf.keras.layers.LSTM(units=64), tf.keras.layers.Dense(num_class, activation='softmax') ]) return model model = create_model(vocab_size=10) model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
モデルのサマリを出力。
pythonコード(続き)
model.summary()
出力結果。
Layer (type) Output Shape Param # ================================================================= embedding (Embedding) (None, None, 80) 800 _________________________________________________________________ lstm (LSTM) (None, 64) 37120 _________________________________________________________________ dense (Dense) (None, 3) 195 ================================================================= Total params: 38,115 Trainable params: 38,115 Non-trainable params: 0 _________________________________________________________________
checkpointデータを出力するように指定してトレーニング。
pythonコード(続き)
# コールバック定義 checkpoint_path = '/tmp/checkpoint_test/LSTM_test' cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path, save_weights_only=True, verbose=1) model.fit(train_images, categorical_labels, epochs=50, verbose=1, callbacks = [cp_callback])
(おまけ)50エポック後のモデルによるsoftmax出力結果
Input: [1, 2, 3, 4] (正解ラベル"0")
tf.Tensor([[0.9474195 0.04316333 0.00941716]], shape=(1, 3), dtype=float32)
=> ラベル"0"に対しては94.7%の確率
Input: [3, 4, 5, 6] (正解ラベル"1")
tf.Tensor([[0.04550831 0.9349791 0.01951269]], shape=(1, 3), dtype=float32)
=>ラベル"1"に対しては93.5%の確率
Input: [5, 6, 7, 8] (正解ラベル"2" )
tf.Tensor([[0.01436339 0.02530523 0.9603314 ]], shape=(1, 3), dtype=float32)
=>ラベル"2"に対しては96.0%の確率
checkpointのパラメータを出力。
出力方法は下記ご参照。
TensorFlowメモ(checkpointの中身を確認する) - ichou1のブログ
レイヤ1 (Embedding)
----------- [param name]: layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (10, 80)
レイヤ2 (LSTM)
----------- [param name]: layer_with_weights-1/cell/bias/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (256,) ----------- [param name]: layer_with_weights-1/cell/kernel/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (80, 256) ----------- [param name]: layer_with_weights-1/cell/recurrent_kernel/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (64, 256)
パラメータの合計は「37120」(= (80 * 64 * 4) + (64 * 64 * 4) + (64 * 4))
# input gate weight(input) : 80row * 64col = 5120 --> self.kernel_i weight(state) : 64row * 64col = 4096 --> self.recurrent_kernel_i bias : 64 --> self.bias_i # forget gate weight(input) : 80row * 64col = 5120 --> self.kernel_f weight(state) : 64row * 64col = 4096 --> self.recurrent_kernel_f bias : 64 --> self.bias_f # input modulation gate weight(input) : 80row * 64col = 5120 --> self.kernel_z weight(state) : 64row * 64col = 4096 --> self.recurrent_kernel_z bias : 64 --> self.bias_z # output gate weight(input) : 80row * 64col = 5120 --> self.kernel_o weight(state) : 64row * 64col = 4096 --> self.recurrent_kernel_o bias : 64 --> self.bias_o
LSTMのパラメータについては下記ご参照。
TensorFlowメモ(RNNその2) - ichou1のブログ
レイヤ3 (Dense)
----------- [param name]: layer_with_weights-2/bias/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (3,) ----------- [param name]: layer_with_weights-2/kernel/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (64, 3) -----------
「tf.compat.v1.keras.layers.CuDNNLSTM」を使った実装
「tf.compat.v1.keras.layers.CuDNNLSTM」を使って、前述のモデル定義を置き換えてみる。
tf.compat.v1.keras.layers.CuDNNLSTM | TensorFlow Core v2.2.0
pythonコード
# Define Model def create_model(vocab_size=10): model = tf.keras.Sequential([ tf.keras.layers.Embedding(vocab_size, 80), tf.compat.v1.keras.layers.CuDNNLSTM(units=64), tf.keras.layers.Dense(num_class, activation='softmax') ]) return model
モデルサマリの出力結果
_________________________________________________________________ Layer (type) Output Shape Param # ================================================================= embedding (Embedding) (None, None, 80) 800 _________________________________________________________________ cu_dnnlstm (CuDNNLSTM) (None, 64) 37376 _________________________________________________________________ dense (Dense) (None, 3) 195 ================================================================= Total params: 38,371 Trainable params: 38,371 Non-trainable params: 0
checkpointのパラメータを出力。
レイヤ2 (CudnnLSTM)
----------- [param name]: layer_with_weights-1/bias/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (512,) ----------- [param name]: layer_with_weights-1/kernel/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (80, 256) ----------- [param name]: layer_with_weights-1/recurrent_kernel/.ATTRIBUTES/VARIABLE_VALUE [param shape]: (64, 256) -----------
bias部分のパラメータ数が異なっている。
「LSTM」の場合は(256, )
ユーザガイドを見ると、"double bias"がデフォルトになっている。
cudnnRNNBiasMode_t is an enumerated type used to specify the number of bias vectors for RNN functions.
Values
CUDNN_RNN_DOUBLE_BIAS
Applies RNN cell formulas that use two bias vectors.
また、DeepSpeechのチェックポイントデータで見たような構成とはなっていない。
DeepSpeechメモ(その2) - ichou1のブログ
「tensorflow.contrib」モジュールを使ったモデルのチェックポイントデータでは、「cudnn_lstm/opaque_kernel」として、パラメータが1つにまとまって出力される。
----------- [param name]: cudnn_lstm/opaque_kernel [param shape]: (33570816,)
Fine-Tuningでトレーニング済みのパラメータを使いたい場合は、該当するパラメータごとに切り出しの上、ロードする必要がありそう。