ichou1のブログ

主に音声認識、時々、データ分析のことを書く

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でトレーニング済みのパラメータを使いたい場合は、該当するパラメータごとに切り出しの上、ロードする必要がありそう。