Python: PyTorch の RNN を検算してみる

今回は、PyTorch の RNN (Recurrent Neural Network) が内部的にどんな処理をしているのか確認してみる。なお、ここでいう RNN は、再起的な構造をもったニューラルネットワークの総称ではなく、いわゆる古典的な Simple RNN を指している。

これを書いている人は、ニューラルネットワークが何もわからないので、再帰的な構造があったりすると尚更わからなくなる。そこで、中身について知っておきたいと考えたのがモチベーションになっている。

使った環境は次のとおり。

$ sw_vers 
ProductName:    macOS
ProductVersion: 11.5.2
BuildVersion:   20G95
$ python -V         
Python 3.9.6
$ pip list | grep torch       
torch                    1.9.0

下準備

まずは PyTorch をインストールしておく。

$ pip install torch

そして、Python のインタプリタを起動する。

$ python

起動できたら PyTorch のモジュールをインポートしておく。

>>> import torch
>>> from torch import nn

モデルを用意する

PyTorch には nn モジュール以下に RNN というクラスがある。このクラスが、ミニバッチに対応した Simple RNN を実装している。このクラスは、最低限 input_size と hidden_size という引数を指定すればインスタンス化できる。

>>> input_dim = 3  # モデルの入力ベクトルの次元数
>>> hidden_dim = 4  # モデルの出力ベクトルの次元数
>>> model = nn.RNN(input_size=input_dim, hidden_size=hidden_dim)

input_size はモデルに入力するデータの次元数で、hidden_size はモデルが出力するデータの次元数になる。 Simple RNN が出力するデータには隠れ状態 (Hidden State) ベクトルという名前がついていて、それが引数の名前に反映されている。

インスタンス化できたら、モデルに含まれるパラメータを確認してみよう。これは何も学習していない状態の初期値だけど、ダミーのデータを使って検算する分にはそれで問題ない。

>>> from pprint import pprint
>>> pprint(list(model.named_parameters()))
[('weight_ih_l0',
  Parameter containing:
tensor([[ 0.4349,  0.2858, -0.3802],
        [ 0.3035,  0.4744, -0.4774],
        [ 0.4553,  0.1563, -0.0048],
        [-0.4107, -0.4734,  0.3651]], requires_grad=True)),
 ('weight_hh_l0',
  Parameter containing:
tensor([[-0.4045,  0.4994, -0.3950,  0.3627],
        [-0.4304,  0.2032,  0.2878,  0.0923],
        [ 0.0641, -0.0405, -0.2965, -0.3422],
        [ 0.3323, -0.2716, -0.1380,  0.2079]], requires_grad=True)),
 ('bias_ih_l0',
  Parameter containing:
tensor([-0.2928,  0.2330,  0.1649, -0.2679], requires_grad=True)),
 ('bias_hh_l0',
  Parameter containing:
tensor([-0.0034, -0.0927,  0.0520, -0.0646], requires_grad=True))]

モデルには 4 つの名前つきパラメータが確認できる。これらのパラメータが何を意味しているかは、以下のドキュメントをみるとわかる。

pytorch.org

上記には、RNN の具体的な計算式が記載されている。

$\displaystyle{ h_t = \tanh(W_{ih} x_t + b_{ih} + W_{hh} h_{(t-1)} + b_{hh}) }$

ここで、 $x_t$ は入力となる系列データにおいて t 番目 (時点) の要素を表していて、 $h_t$ は t 番目の隠れ状態ベクトルになる。 $h_{(t-1)}$ は t - 1 番目の隠れ状態ベクトルなので、1 つ前の状態の出力を入力として使っていることがわかる。

それ以外は、先ほどのパラメータと次のように対応している。

$W_{ih}$
- weight_ih_l0
$W_{hh}$
- weight_hh_l0
$b_{ih}$
- bias_ih_l0
$b_{hh}$
- bias_hh_l0

ダミーデータを用意する

式がわかったところで、検算するための出力を適当に用意したダミーデータを使って得よう。次のようにランダムな入力データを用意する。

>>> T = 5  # 入力する系列データの長さ
>>> batch_size = 2  # 一度に処理するデータの数
>>> X = torch.randn(T, batch_size, input_dim)  # ダミーの入力データ
>>> X.shape
torch.Size([5, 2, 3])

上記のダミーデータをモデルに入力として与える。すると、タプルで 2 つの返り値が得られる。

>>> H, hn = model(X)

このうち、タプルの最初の要素は各時点 (0 ~ T) での隠れ状態ベクトルが入っている。つまり、X[0] に対応した隠れ状態ベクトルが H[0] で、X[1] に対応した隠れ状態ベクトルが H[1] で...ということ。

>>> H.shape
torch.Size([5, 2, 4])
>>> H
tensor([[[-0.0096,  0.3380,  0.4147, -0.5187],
         [-0.5797,  0.0438, -0.3449, -0.0454]],

        [[-0.3769,  0.5505,  0.1542, -0.6927],
         [ 0.1021,  0.4838,  0.0174, -0.5226]],

        [[ 0.5723,  0.8306,  0.5878, -0.9012],
         [-0.5423, -0.3730,  0.1816,  0.0130]],

        [[-0.2641,  0.0466,  0.7226, -0.6048],
         [-0.6680, -0.4764,  0.2837,  0.2118]],

        [[-0.8623, -0.3724, -0.4284,  0.2948],
         [-0.2464,  0.4500, -0.4194, -0.1977]]], grad_fn=<StackBackward>)

そして、タプルで 2 番目に返ってきた値は最後の時点 (T) での隠れ状態ベクトルになる。ようするに、上記の最後尾と同じもの。

>>> hn
tensor([[[-0.8623, -0.3724, -0.4284,  0.2948],
         [-0.2464,  0.4500, -0.4194, -0.1977]]], grad_fn=<StackBackward>)
>>> H[-1]
tensor([[-0.8623, -0.3724, -0.4284,  0.2948],
        [-0.2464,  0.4500, -0.4194, -0.1977]], grad_fn=<SelectBackward>)

検算する

次に、実際の検算に入る。まずは、次のようにして各パラメータの Tensor オブジェクトを得る。

>>> model_weights = {name: param.data for name, param
...                  in model.named_parameters()}
>>> 
>>> W_ih = model_weights['weight_ih_l0']
>>> W_hh = model_weights['weight_hh_l0']
>>> b_ih = model_weights['bias_ih_l0']
>>> b_hh = model_weights['bias_hh_l0']

まずは系列データの一番最初の t = 0 で X[0] に対応する隠れ状態ベクトルから求める。ターゲットはこれ。

>>> H[0]
tensor([[-0.0096,  0.3380,  0.4147, -0.5187],
        [-0.5797,  0.0438, -0.3449, -0.0454]], grad_fn=<SelectBackward>)

やることは単純で、先ほどの式を PyTorch で表現すれば良い。なお、t = 0 の時点では $h_{(t-1)}$ がないので、その項は消える。

>>> torch.tanh(torch.matmul(W_ih, X[0].T).T + b_ih + b_hh)
tensor([[-0.0096,  0.3380,  0.4147, -0.5187],
        [-0.5797,  0.0438, -0.3449, -0.0454]])

Tensor の値が一致していることがわかる。

次は t = 1 で X[1] に対応する隠れ状態ベクトルを求める。ターゲットは以下。

>>> H[1]
tensor([[-0.3769,  0.5505,  0.1542, -0.6927],
        [ 0.1021,  0.4838,  0.0174, -0.5226]], grad_fn=<SelectBackward>)

t = 1 では $h_{(t-1)}$ が $h_0$ になる。とはいえ項が増えるだけで、やることは先ほどと変わらない。

>>> torch.tanh(torch.matmul(W_ih, X[1].T).T + b_ih + torch.matmul(W_hh, H[0].T).T + b_hh)
tensor([[-0.3769,  0.5505,  0.1542, -0.6927],
        [ 0.1021,  0.4838,  0.0174, -0.5226]], grad_fn=<TanhBackward>)

こちらも値が一致している。

あとは添字が増えるだけなので省略する。

初期 (`t = 0`) の隠れ状態ベクトルを渡す場合

先ほどの例では、初期 (t = 0) のときに $h_{(t-1)}$ に相当する隠れ状態ベクトルが存在しなかった。これは自分で用意して渡すこともできるので、その場合の挙動も確認しておこう。

次のようにしてランダムな値で初期の隠れ状態ベクトルを h0 として用意する。なお、先頭の次元は Simple RNN を重ねる段数を表している。というのも、(総称としての) RNN は縦に積み重ねることで性能向上が望める場合があるらしい ¹。 PyTorch のRNN も、インスタンス化するときに num_layers という引数で重ねる数が指定できる。なお、デフォルト値は 1 になっている。

>>> rnn_layers = 1  # Simple RNN を重ねる数 (num_layers の値)
>>> h0 = torch.randn(rnn_layers, batch_size, hidden_dim)
>>> h0.shape
torch.Size([1, 2, 4])

初期の隠れ状態ベクトルをモデルに渡すには、次のように 2 番目の引数として渡せば良い。

>>> H, hn = model(X, h0)

先ほどと同じように検算してみよう。

>>> H[0]
tensor([[-0.1925,  0.6594, -0.2041, -0.4893],
        [ 0.5740,  0.8465, -0.5979, -0.8112]], grad_fn=<SelectBackward>)

といっても、最初の $h_{(t-1)}$ として h0 を使うだけ。

>>> torch.tanh(torch.matmul(W_ih, X[0].T).T + b_ih + torch.matmul(W_hh, h0[0].T).T + b_hh)
tensor([[-0.1925,  0.6594, -0.2041, -0.4893],
        [ 0.5740,  0.8465, -0.5979, -0.8112]])

残りは変わらない。

>>> H[1]
tensor([[ 0.0929,  0.5283,  0.2951, -0.7210],
        [-0.1403,  0.0510,  0.3764, -0.4921]], grad_fn=<SelectBackward>)
>>> torch.tanh(torch.matmul(W_ih, X[1].T).T + b_ih + torch.matmul(W_hh, H[0].T).T + b_hh)
tensor([[ 0.0929,  0.5283,  0.2951, -0.7210],
        [-0.1403,  0.0510,  0.3764, -0.4921]], grad_fn=<TanhBackward>)

Simple RNN を重ねた場合

先ほど述べたとおり RNN は層を重ねることで性能向上が望める場合がある。その場合についても確認しておく。

まずは RNN を 2 層重ねたモデルを用意する。

>>> rnn_layers = 2
>>> model = nn.RNN(input_size=input_dim, hidden_size=hidden_dim, num_layers=rnn_layers)

次のようにモデルのパラメータが増えている。具体的には名前の末尾が l0 になったものと l1 になったものがある。これはつまりl0 の上に l1 が重なっていることを示す。

>>> pprint(list(model.named_parameters()))
[('weight_ih_l0',
  Parameter containing:
tensor([[-0.3591,  0.0948, -0.0500],
        [ 0.1963, -0.1717, -0.3551],
        [ 0.0313,  0.0495, -0.0878],
        [ 0.3109,  0.3728,  0.2577]], requires_grad=True)),
 ('weight_hh_l0',
  Parameter containing:
tensor([[-0.3050, -0.0269,  0.1772,  0.0081],
        [-0.0770,  0.3563, -0.1209,  0.0126],
        [-0.3534,  0.0264,  0.2649,  0.2235],
        [ 0.3338, -0.0708,  0.4314, -0.0149]], requires_grad=True)),
 ('bias_ih_l0',
  Parameter containing:
tensor([ 0.3767,  0.3653, -0.1024,  0.3425], requires_grad=True)),
 ('bias_hh_l0',
  Parameter containing:
tensor([-0.1083, -0.1802, -0.2972,  0.1099], requires_grad=True)),
 ('weight_ih_l1',
  Parameter containing:
tensor([[ 0.2279, -0.4886,  0.4573,  0.2441],
        [-0.0949, -0.2300,  0.1320, -0.2643],
        [ 0.0720,  0.4727,  0.2005, -0.0784],
        [-0.0784,  0.3208,  0.4977, -0.0190]], requires_grad=True)),
 ('weight_hh_l1',
  Parameter containing:
tensor([[-0.0565,  0.1433,  0.0810,  0.1619],
        [ 0.2734,  0.3270, -0.2813,  0.1076],
        [ 0.2989,  0.0412, -0.1173,  0.1614],
        [-0.0805, -0.1851, -0.1254,  0.0713]], requires_grad=True)),
 ('bias_ih_l1',
  Parameter containing:
tensor([-0.3898, -0.1349, -0.2269, -0.1637], requires_grad=True)),
 ('bias_hh_l1',
  Parameter containing:
tensor([ 0.4969,  0.3327,  0.4548, -0.3809], requires_grad=True))]

それぞれのパラメータの重みを取得しておく。

>>> model_weights = {name: param.data for name, param
...                  in model.named_parameters()}
>>> 
>>> W_ih_l0 = model_weights['weight_ih_l0']
>>> W_hh_l0 = model_weights['weight_hh_l0']
>>> b_ih_l0 = model_weights['bias_ih_l0']
>>> b_hh_l0 = model_weights['bias_hh_l0']
>>> 
>>> W_ih_l1 = model_weights['weight_ih_l1']
>>> W_hh_l1 = model_weights['weight_hh_l1']
>>> b_ih_l1 = model_weights['bias_ih_l1']
>>> b_hh_l1 = model_weights['bias_hh_l1']

入力のダミーデータはそのままに、モデルからあらためて隠れ状態ベクトルを取得する。

>>> H, hn = model(X)

初期状態 (t = 0) をターゲットにする。

>>> H[0]
tensor([[-0.1115, -0.0662,  0.2981, -0.5452],
        [ 0.0896,  0.0750,  0.2003, -0.6533]], grad_fn=<SelectBackward>)

まずは、これまでの要領で隠れ状態ベクトルを得る。ただし、これはあくまで 1 層目の出力にすぎない。使っているパラメータの名前も末尾が _l0 になっている。

>>> h0_l0 = torch.tanh(torch.matmul(W_ih_l0, X[0].T).T + b_ih_l0 + b_hh_l0)
>>> h0_l0
tensor([[ 0.0724,  0.3836, -0.3525,  0.4635],
        [ 0.6664,  0.0096, -0.3751,  0.0292]])

続いて、1 層目の出力を 2 層目に入力して計算する。

>>> torch.tanh(torch.matmul(W_ih_l1, h0_l0.T).T + b_ih_l1 + b_hh_l1)
tensor([[-0.1115, -0.0662,  0.2981, -0.5452],
        [ 0.0896,  0.0750,  0.2003, -0.6533]])

これで値が一致した。

そんなかんじで。

参考書籍

ゼロから作るDeep Learning ❷ ―自然言語処理編

作者:斎藤康毅
オライリー・ジャパン

Amazon

詳しくは「ゼロから作るDeep Learning ❷ ―自然言語処理編」を参照のこと↩

CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: PyTorch の RNN を検算してみる

もくじ

下準備

モデルを用意する

ダミーデータを用意する

検算する

初期 (`t = 0`) の隠れ状態ベクトルを渡す場合

Simple RNN を重ねた場合

参考書籍

もくじ

下準備

モデルを用意する

ダミーデータを用意する

検算する

初期 (t = 0) の隠れ状態ベクトルを渡す場合

Simple RNN を重ねた場合

参考書籍

初期 (`t = 0`) の隠れ状態ベクトルを渡す場合