音声認識エンジンの深層学習フレームワークをChainerからPyTorchに移行しました
※こちらの記事は、2020年8月11日にRetrieva TECH BLOGにて掲載された記事を再掲載したものとなります。
音声認識チームのソフトウェアエンジニアの西岡です。
弊社では音声認識エンジンを開発しており、これまでChainerを使って音声認識モデルの訓練・精度評価をおこなってきましたが、Chainer v7を最後に開発がストップすることが発表されたため、今回ChainerからPyTorchへの移行をおこないました。
本記事では、移行にあたってぶつかった問題や工夫した点について紹介します。
音声認識エンジンのシステム構成
弊社で開発している音声認識エンジンは、音声認識モデルを訓練する部分と訓練済みのモデルを使って推論する部分でシステムが分かれています。以降、簡単のためそれぞれ「訓練側」「推論側」と呼称することにします。
訓練側は基本的にPythonで実装されており、Chainerを使用して音声認識モデルを訓練し、ディスクに保存します。
推論側ではディスクで保存されたモデルを読み込んで推論をおこないます。 推論側は弊社の製品として提供され、GPUの無い環境でも高速に推論をおこなうことが期待されます。そのため、Chainerのforward処理と同じ結果を出力する推論処理がC++で独自に実装されています。
PyTorch移行にあたってぶつかった問題
本節では、訓練側で使用されていたフレームワークをChainerからPyTorchに移行するにあたってぶつかった問題と、その問題を解決するべく工夫した点を述べていきます。
問題1: LSTMが受け取る入力の仕様の違い
Chainerの NStepLSTM では 系列長x入力サイズ の chainer.Variable が入ったバッチサイズ分の長さのリストを入力として受け付ける一方、PyTorchの LSTM では 系列長xバッチサイズx入力サイズ の torch.Tensor を入力として受け付けます。
系列長の異なるサンプルを含んだバッチをLSTMに入力する場合、Chainerでは元の系列長のままリストに持たせればよいのに対し、PyTorchでは全サンプルで同じ系列長になるようpaddingする必要があります。 結果、PyTorchのLSTMではpaddingした部分のタイムステップも含めて計算されてしまい、推論結果がChainerのLSTMのそれとは変わってしまいます。
この問題は、入力のバッチを PackedSequence というオブジェクトに変換することで解決できます。
>>> import torch
>>> from torch.nn.utils.rnn import pack_sequence
>>> input_size = 3
>>> input_lengths = [5, 4, 3]
>>> xs = [torch.randn(length, input_size) for length in input_lengths]
# `サンプル目、系列長5 x 入力サイズ3 のTensor
>>> xs[0]
>>> tensor([[-0.3519, 0.0197, -0.7828],
[-0.2282, 0.3611, -0.3079],
[-0.1541, 0.8168, -0.4285],
[-0.8118, -2.8849, -0.1072],
[-1.0132, 0.7138, -0.3199]])
# 2サンプル目、系列長4 x 入力サイズ3 のTensor
>>> xs[1]
tensor([[-0.6646, -1.4412, 0.7494],
[ 1.4010, -1.3183, -0.1103],
[-0.2029, -0.8214, -1.3497],
[-1.7019, -0.2867, -2.0249]])
# 3サンプル目、系列長3 x 入力サイズ3 のTensor
>>> xs[2]
tensor([[-0.4645, 0.9132, -2.0173],
[-0.7927, -0.4681, 0.5647],
[ 0.2981, -0.9253, 0.7447]])
>>> packed_input = pack_sequence(xs)
>>> packed_input
PackedSequence(data=tensor([[-0.3519, 0.0197, -0.7828],
[-0.6646, -1.4412, 0.7494],
[-0.4645, 0.9132, -2.0173],
[-0.2282, 0.3611, -0.3079],
[ 1.4010, -1.3183, -0.1103],
[-0.7927, -0.4681, 0.5647],
[-0.1541, 0.8168, -0.4285],
[-0.2029, -0.8214, -1.3497],
[ 0.2981, -0.9253, 0.7447],
[-0.8118, -2.8849, -0.1072],
[-1.7019, -0.2867, -2.0249],
[-1.0132, 0.7138, -0.3199]]), batch_sizes=tensor([3, 3, 3, 2, 1]), sorted_indices=None, unsorted_indices=None)
系列長x入力サイズ の torch.Tensor が入ったバッチサイズ分の長さのリストを torch.nn.utils.rnn.pack_sequence に渡すと PackedSequence オブジェクトが返されます。 PackedSequence.data は 全サンプルの系列長の合計x入力サイズ の2次元の配列になっており、1サンプル目の1タイムステップ目、2サンプル目の1タイムステップ目、3サンプル目の1タイムステップ目、1サンプル目の2タイムステップ目、というようにサンプル順xタイムステップ順で値が入っています。
この PackedSequence オブジェクトを LSTM に入力して推論させてみましょう。 PackedSequence が渡された時は返り値も PackedSequence なので、 pad_packed_sequence で Tensor に変換します。
>>> from torch.nn.utils.rnn import pad_packed_sequence
>>> hidden_size = 5
>>> lstm = torch.nn.LSTM(input_size, hidden_size)
>>> lstm.eval()
>>> with torch.no_grad():
... packed_output, hidden = lstm(packed_input)
...
>>> padded_output, _ = pad_packed_sequence(packed_output)
>>> padded_output
tensor([[[-0.1149, -0.1032, -0.0465, -0.0451, 0.0676],
[ 0.0701, 0.0203, -0.0414, -0.1424, -0.0077],
[-0.2667, -0.1315, -0.0541, 0.0727, 0.1405]],
[[-0.1877, -0.1722, -0.0720, -0.0658, -0.0027],
[ 0.2580, 0.1275, 0.0995, -0.1037, 0.0769],
[-0.1353, -0.1261, -0.1219, -0.1167, 0.0199]],
[[-0.2803, -0.2288, -0.0783, -0.0237, -0.0639],
[ 0.0812, 0.0006, 0.0487, -0.0917, 0.2095],
[ 0.0179, -0.0280, -0.0915, -0.1580, -0.0517]],
[[-0.0688, 0.0116, -0.0432, -0.0652, 0.1299],
[-0.1218, -0.0766, -0.0935, -0.0129, 0.2719],
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
[[-0.1660, -0.0982, -0.0831, -0.0464, 0.0692],
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]])
2サンプル目の5ステップ目と3サンプル目の4,5ステップ目はpadされた値が入っていることが確認できます。
一方、単純に pad_sequence でpaddingした Tensor をLSTMに入力した結果はこちらです。
>>> from torch.nn.utils.rnn import pad_sequence
>>> padded_input = pad_sequence(xs)
>>> with torch.no_grad():
... output, hidden = lstm(padded_input)
...
>>> output
tensor([[[-0.1149, -0.1032, -0.0465, -0.0451, 0.0676],
[ 0.0701, 0.0203, -0.0414, -0.1424, -0.0077],
[-0.2667, -0.1315, -0.0541, 0.0727, 0.1405]],
[[-0.1877, -0.1722, -0.0720, -0.0658, -0.0027],
[ 0.2580, 0.1275, 0.0995, -0.1037, 0.0769],
[-0.1353, -0.1261, -0.1219, -0.1167, 0.0199]],
[[-0.2803, -0.2288, -0.0783, -0.0237, -0.0639],
[ 0.0812, 0.0006, 0.0487, -0.0917, 0.2095],
[ 0.0179, -0.0280, -0.0915, -0.1580, -0.0517]],
[[-0.0688, 0.0116, -0.0432, -0.0652, 0.1299],
[-0.1218, -0.0766, -0.0935, -0.0129, 0.2719],
[-0.0217, -0.0961, -0.0830, -0.1611, -0.0930]],
[[-0.1660, -0.0982, -0.0831, -0.0464, 0.0692],
[-0.1446, -0.1527, -0.0931, -0.0967, 0.1841],
[-0.0509, -0.1312, -0.0851, -0.1512, -0.1079]]])
padされた箇所もLSTMで推論されてしまっていることが確認できます。
最初は何も考えずに pad_sequence でpaddingしていたため、単体テストでChainerと推論結果が一致せず苦労しましたが、そもそも LSTMのドキュメントに pack_padded_sequence か pack_sequence を使えと書いてあるので、ドキュメントをちゃんと読めという話でした。
問題2: PyTorchではモデルをhdf5形式で保存できない
システム構成の節でも述べましたが推論側はC++で実装されているため、C++で読み込める形式でモデルを保存する必要があります。
Chainerで訓練したモデルは chainer.serializers.save_hdf5を使ってhdf5形式で保存しており、推論側はそのhdf5形式で保存されたモデルを読み込んで推論に使用しています。
しかし、PyTorchではモデルをhdf5形式で保存する仕組みが提供されていません。
PyTorchにはTorchScriptというNNモジュールをPython処理系とは独立した形式に変換する仕組みがあり、モデルをTorchScript形式で保存することでPyTorchのC++ APIを使って推論することが可能ですが、推論側がPyTorchのC++ APIにがっつり依存してしまうことになる上、そのための開発工数も必要です。
そこで、PyTorchのNNモジュールをhdf5形式でシリアライズ・デシリアライズ（以下簡単のためSerDeと呼称します）する処理を自作することにしました。これにより、訓練側だけの開発だけで済み、推論側に開発は必要ありません。
ただし、新しいモジュールを音声認識モデルで使用するとhdf5のSerDeを実装しなければならないというトレードオフがあるため、今後モデルの改善が活発になればPyTorch C++ APIに依存してしまうというのも選択肢の1つとしてはあり得ますが、今回はPyTorchへの移行を優先するために前者を選択しました。
hdf5とはモデルを階層的な構造で保存することができるフォーマットです。
例えば、Chainerで多層パーセプトロンを実装してみます。
import chainer
import chainer.functions as F
import chainer.links as L
class MLP(chainer.Chain):
def __init__(self, input_size, hidden_size=10, output_size=3):
super().__init__()
with self.init_scope():
self.fc1 = L.Linear(input_size, hidden_size)
self.fc2 = L.Linear(hidden_size, hidden_size)
self.fc3 = L.Linear(hidden_size, output_size)
def forward(self, x):
h = F.relu(self.fc1(x))
h = F.relu(self.fc2(h))
h = self.fc3(h)
return h
このモデルをインスタンス化し、パラメータの一覧を表示してみます。
>>> input_size = 5
>>> model = MLP(input_size)
>>> for name, param in model.namedparams():
... print('name: {}, param: {}'.format(name, param.shape))
...
name: /fc1/W, param: (10, 5)
name: /fc1/b, param: (10,)
name: /fc2/W, param: (10, 10)
name: /fc2/b, param: (10,)
name: /fc3/W, param: (3, 10)
name: /fc3/b, param: (3,)
nn.Linear は W と b の2つのパラメータを持ち、それが変数名と / で結合された名前で保持されていることがわかります。
このモデルを chainer.serializers.save_hdf5 で保存し、HDFViewで表示させてみましょう。
>>> import chainer
>>> import h5py
>>> filepath = 'mlp.hdf5'
>>> chainer.serializers.save_hdf5(filepath, model)
/fc1/W と /fc1/b が同じ階層に位置しており、階層的にデータを保持するフォーマットであるということが視覚的にわかりやすくなります。
次に、hdf5ファイルからパラメータを取得する処理を書いてみます。
>>> import h5py
>>> with h5py.File(filepath, 'r') as f:
... data = f['/fc2/W']
... print(type(data))
... print(data.shape)
...
<class 'h5py._hl.dataset.Dataset'>
(10, 10)
hdf5ファイルをopenし、 ファイルオブジェクトに対し '/fc2/W' をキーとしてインデックスアクセスすると、対応するパラメータが取得できます。
では、PyTorchでも同様の多層パーセプトロンを実装してみます。
import torch.nn as nn
import torch.nn.functional as F
class MLP(nn.Module):
def __init__(self, input_size, hidden_size=10, output_size=3):
super().__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.fc3 = nn.Linear(hidden_size, output_size)
def forward(self, x):
h = F.relu(self.fc1(x))
h = F.relu(self.fc2(h))
h = self.fc3(h)
return h
このモデルの持つパラメータを出力させてみます。
>>> input_size = 5
>>> model = MLP(input_size)
>>> for name, param in model.named_parameters():
... print('name: {}, param: {}'.format(name, param.shape))
...
name: fc1.weight, param: torch.Size([10, 5])
name: fc1.bias, param: torch.Size([10])
name: fc2.weight, param: torch.Size([10, 10])
name: fc2.bias, param: torch.Size([10])
name: fc3.weight, param: torch.Size([3, 10])
name: fc3.bias, param: torch.Size([3])
nn.Linear の持つパラメータ名が weight と bias になっている点、変数名とパラメータ名が . で結合された形式になっている点はChainerとは異なりますが、シリアライズ時は nn.Linear の fc1.weight をhdf5ファイルの /fc1/W パスに保存、デシリアライズ時は /fc1/W パスを指定してhdf5ファイルからパラメータを取得し nn.Linear の fc1.weight にセットするような処理を書けば、hdf5のSerDeが実現できそうです。
そこで、下のような処理を実装しました。まずはシリアライズです。
import h5py
import torch
import torch.nn as nn
def serialize_linear(f, module, base_key):
param_generator = module.parameters()
weight = next(param_generator)
key = '/'.join([base_key, 'W'])
f.create_dataset(key, data=weight.detach())
bias = next(param_generator)
key = '/'.join([base_key, 'b'])
f.create_dataset(key, data=bias.detach())
def serialize(f, module, base_key):
if isinstance(module, nn.Linear):
serialize_linear(f, module, base_key)
else:
for name, child in module.named_children():
serialize(f, child, '/'.join([base_key, name]))
def save_hdf5(filepath, module):
with h5py.File(filepath, 'w') as f:
serialize(f, module, '')
chainer.serialziers.save_hdf5 と同じシグネチャを持つ save_hdf5 関数でPyTorchのNNモジュールをhdf5で保存します。 save_hdf5 関数はhdf5ファイルをopenした後 serialize 関数を呼び出します。 serialize 関数はモジュールを深さ優先で探索し、 nn.Linear のインスタンスを見つけた場合 serialize_linear 関数を呼び出し、create_dataset メソッドにキーを指定して重みをhdf5ファイルに保存します。
次にデシリアライズです。
import h5py
import torch
import torch.nn as nn
def deserialize_linear(f, module, base_key):
param_generator = module.parameters()
weight = next(param_generator)
key = '/'.join([base_key, 'W'])
dataset = f[key]
weight.data = torch.Tensor(dataset)
bias = next(param_generator)
key = '/'.join([base_key, 'b'])
dataset = f[key]
bias.data = torch.Tensor(dataset)
def deserialize(f, module, base_key):
if isinstance(module, nn.Linear):
deserialize_linear(f, module, base_key)
else:
for name, child in module.named_children():
deserialize(f, child, '/'.join([base_key, name]))
def load_hdf5(filepath, module):
with h5py.File(filepath, 'r') as f:
deserialize(f, module, '')
シリアライズと非常に似た実装になっており、 deserialize_linear 関数でhdf5ファイルから重みを読み込んでNNモジュールにセットしています。
上の例では nn.Linear のSerDeをおこなっていますが、 serialize 関数と deserialize 関数内の if isinstance(module, nn.Linear) の分岐を別のNNモジュールについても追加し、対応する serialize_xxx 関数と deserialize_xxx 関数を実装することで、他のNNモジュールもhdf5のSerDeをおこなうことができます。例えば nn.LSTM であれば、 if isinstance(module, nn.LSTM) の分岐を追加し serialize_lstm 関数と deserialize_lstm 関数を実装することで対応できます。
しかし、LSTMについては一筋縄ではいきませんでした。それについては次の節で説明したいと思います。
問題3: LSTMのパラメータの持ち方がChainerとPyTorchで異なる
本項で記載している内容はLSTMの内部構造を知っていると理解が容易になると思います。LSTMの内部構造はUnderstanding LSTM Networks -- colah's blogにわかりやすく記載されているので一読することをお勧めします。
上で見てきたように、Linearのパラメータの持ち方はChainerとPyTorchであまり変わりませんでしたが、 LSTMは大きく異なります。
2層・bidirectionalなLSTMのパラメータを表示させてみます。まずはChainerの NStepBiLSTM です。
>>> import chainer.links as L
>>> n_layers, in_size, out_size, dropout = 2, 3, 5, 0.0
>>> lstm = L.NStepBiLSTM(n_layers, in_size, out_size, dropout)
>>> for name, param in lstm.namedparams():
... print(name, param.shape)
...
/0/b0 (5,)
/0/b1 (5,)
/0/b2 (5,)
/0/b3 (5,)
/0/b4 (5,)
/0/b5 (5,)
/0/b6 (5,)
/0/b7 (5,)
/0/w0 (5, 3)
/0/w1 (5, 3)
/0/w2 (5, 3)
/0/w3 (5, 3)
/0/w4 (5, 5)
/0/w5 (5, 5)
/0/w6 (5, 5)
/0/w7 (5, 5)
/1/b0 (5,)
/1/b1 (5,)
/1/b2 (5,)
/1/b3 (5,)
/1/b4 (5,)
/1/b5 (5,)
/1/b6 (5,)
/1/b7 (5,)
/1/w0 (5, 3)
/1/w1 (5, 3)
/1/w2 (5, 3)
/1/w3 (5, 3)
/1/w4 (5, 5)
/1/w5 (5, 5)
/1/w6 (5, 5)
/1/w7 (5, 5)
/2/b0 (5,)
/2/b1 (5,)
/2/b2 (5,)
/2/b3 (5,)
/2/b4 (5,)
/2/b5 (5,)
/2/b6 (5,)
/2/b7 (5,)
/2/w0 (5, 10)
/2/w1 (5, 10)
/2/w2 (5, 10)
/2/w3 (5, 10)
/2/w4 (5, 5)
/2/w5 (5, 5)
/2/w6 (5, 5)
/2/w7 (5, 5)
/3/b0 (5,)
/3/b1 (5,)
/3/b2 (5,)
/3/b3 (5,)
/3/b4 (5,)
/3/b5 (5,)
/3/b6 (5,)
/3/b7 (5,)
/3/w0 (5, 10)
/3/w1 (5, 10)
/3/w2 (5, 10)
/3/w3 (5, 10)
/3/w4 (5, 5)
/3/w5 (5, 5)
/3/w6 (5, 5)
/3/w7 (5, 5)
/{idx}/{w|b}{0-7} という名前でパラメータが保持されています。
まず、shapeが2次元か1次元かによって w は行列積用の重みで b はバイアスであると容易にわかります。 idx=0,1 のとき w のshapeが (hidden_size, input_size) になっているので、 idx=0,1 はLSTMの1層目、 idx=2,3 は2層目だとわかります。shapeのみからだと idx が偶数と奇数のどちらが順方向用なのか逆方向用なのかはわかりませんが、結論からいうと idx が偶数であれば順方向用・奇数であれば逆方向用のパラメータです。
次に、2層目の w{0-3} のshapeが (hidden_size, hidden_size * 2) 、 w{4-7} のshapeが (hidden_size, hidden_size) になっていることから、 0-3 が入力データ用で 4-7 が隠れ状態（いわゆる h ）用のパラメータであることを示しています。 ただし、shapeのみからだとどの数字がどのゲート用のパラメータなのかはわかりません。
次に、PyTorchの nn.LSTM のパラメータを表示させてみます。
>>> import torch.nn as nn
>>> input_size, hidden_size, num_layers = 3, 5, 2
>>> lstm = nn.LSTM(input_size, hidden_size, num_layers, bidirectional=True)
>>> for name, param in lstm.named_parameters():
... print(name, param.shape)
...
weight_ih_l0 torch.Size([20, 3])
weight_hh_l0 torch.Size([20, 5])
bias_ih_l0 torch.Size([20])
bias_hh_l0 torch.Size([20])
weight_ih_l0_reverse torch.Size([20, 3])
weight_hh_l0_reverse torch.Size([20, 5])
bias_ih_l0_reverse torch.Size([20])
bias_hh_l0_reverse torch.Size([20])
weight_ih_l1 torch.Size([20, 10])
weight_hh_l1 torch.Size([20, 5])
bias_ih_l1 torch.Size([20])
bias_hh_l1 torch.Size([20])
weight_ih_l1_reverse torch.Size([20, 10])
weight_hh_l1_reverse torch.Size([20, 5])
bias_ih_l1_reverse torch.Size([20])
bias_hh_l1_reverse torch.Size([20])
weight は行列積用の重みで bias はバイアス、 ih は入力データ用で hh は隠れ状態用のパラメータ、 l{N} はN層目、 末尾の reverse の有無でunidirectionalかbidirectionalかを表しています。
大きな違いは、Chainerの NStepBiLSTM ではゲートごとにパラメータが分割されて保持されているのに対し、PyTorchの nn.LSTM では結合されて保持されている点です。
前項で説明したように、C++の推論側ではChainerのhdf5形式で保存されたモデルを読み込むため、 推論側で読み込むためにはPyTorchで weight_ih_l0 として保持されているパラメータを4つに分割し /0/w{0-3} のパスに割り当てる必要があります。ただし、どのように分割しどの順番で割り当てるかが自明ではありません。そこで、各フレームワークのLSTMの実装に潜って確認する必要があります。
まず、Chainerの NStepBiLSTM において、LSTMcellのforwardのCPU実装は chainer.functions.rnn.n_step_lstm._lstm にあたります。
def _lstm(x, h, c, w, b):
xw = _stack_weight([w[2], w[0], w[1], w[3]])
hw = _stack_weight([w[6], w[4], w[5], w[7]])
xb = _stack_weight([b[2], b[0], b[1], b[3]])
hb = _stack_weight([b[6], b[4], b[5], b[7]])
lstm_in = linear.linear(x, xw, xb) + linear.linear(h, hw, hb)
c_bar, h_bar = lstm.lstm(c, lstm_in)
return h_bar, c_bar
_stack_weight 関数を使ってゲートごとの重みとバイアスを第0軸方向にスタックさせることで、1回の行列積で全てのゲートの行列積を計算することができます。1層目であればスタック後の xw のshapeは (hidden_size x 4, input_size) になり、 lstm_in のshapeは (batch_size, hidden_size x 4) になります。スタックさせる順番は 2,0,1,3 （隠れパラメータ側は 6,4,5,7 ）です。
行列積の結果は chainer.functions.rnn.lstm.LSTM.forward に渡されます。
def forward(self, inputs):
self.retain_inputs((0, 1))
c_prev, x = inputs
a, i, f, o = _extract_gates(x)
batch = len(x)
L91で、行列積の結果がゲートごとに分割しています。 _extract_gates の操作は _stack_weight の真逆の操作なので、 w[2] と w[6] （と b[2] ・ b[6] 、以降はバイアスは省略します）を使った行列積の結果が a に、 w[0] と w[4] を使った行列積の結果が f に、 w[1] と w[5] を使った行列積の結果が f に、 w[3] と w[7] を使った行列積の結果が o に対応します。
変数名である程度あたりがつけることもできますが、後の処理をもう少し見てみます。
a = xp.tanh(a)
i = _sigmoid(i, xp)
f = _sigmoid(f, xp)
o = _sigmoid(o, xp)
c_next = numpy.empty_like(c_prev)
c_next[:batch] = a * i + f * c_prev[:batch]
h = o * xp.tanh(c_next[:batch])
tanh を通しているので a はセルゲートの出力、 セルゲートとelement-wiseの積を計算しているので i は入力ゲートの出力、前タイムステップからのcell stateとelement-wiseの積を計算しているので f は忘却ゲートの出力、 残った o が出力ゲートの出力であるということがわかります。
このことから、重みのインデックスの 0 と 4 が入力ゲートに、 1 と 5 が忘却ゲートに、 2 と 6 がセルゲートに、 3 と 7 が出力ゲートにそれぞれ対応していることがわかりました。
次に、PyTorchの実装も見てみます。 LSTMCell のCPU実装は src/ATen/native/RNN.cpp の LSTMCell::operator() で実装されています。
const auto gates = params.linear_hh(hx).add_(
pre_compute_input ? input : params.linear_ih(input));
auto chunked_gates = gates.chunk(4, 1);
auto ingate = chunked_gates[0].sigmoid_();
auto forgetgate = chunked_gates[1].sigmoid_();
auto cellgate = chunked_gates[2].tanh_();
auto outgate = chunked_gates[3].sigmoid_();
auto cy = (forgetgate * cx).add_(ingate * cellgate);
auto hy = outgate * cy.tanh();
L686でゲートごとの行列積をまとめて計算し、 L687で第1軸方向に4つに分割しています。 分割した結果、 chunked_gates[0] が入力ゲートの出力に、 chunked_gates[1] が忘却ゲートの出力に、 chunked_gates[2] がセルゲートの出力に、 chunked_gates[3] が出力ゲートの出力に対応するため、重みも第0軸方向に（行列積の結果では第1軸方向でしたが重みは第0軸であることに注意）同様の順番でスタックされているであろうことがわかります。
結論としては、 w0, w1, w2, w3 = weight_ih.chunk(4, 0) , w4, w5, w6, w7 = weight_hh.chunk(4, 0) のようにPyTorchの重みを第0軸方向に4つに分割した結果がそのままChainerの重みに対応することがわかったため、これに基づいて serialize_lstm 関数と deserialize_lstm 関数を実装しています。紙面の都合で実装は省略させていただきます。
まとめ
本記事では、音声認識エンジンの深層学習フレームワークをChainerからPyTorchに移行するにあたってぶつかった問題や工夫した点について述べました。
今後もPyTorchや音声認識について得られた知見を本ブログで共有していければと思います。