書いて理解するPytorchのforwardとbackward

forwardは一言で言えば順伝搬の処理を定義しています。

元々はkerasを利用していましたが、時代はpytorchみたいな雰囲気に呑まれpytorchに移行中です。
ただkerasに比べて複雑に感じる時があります。
今回はforwardを書いていて、「なんだっけこれ」と初心者してしまっておりますので、
pytorch.org
今回はこちらの公式のexampleを実行してみて、理解に努めようと思います。

データセットの用意

必要なライブラリは読み込んでおきます

import numpy as np
import pandas as pd
import torch
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
import matplotlib.pyplot as plt

データセットは今回アヤメを利用しようと思います。

data = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target

バッチサイズは25、隠れ層の次元数も適当に3にしておきます。

iris = load_iris()
data = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target
y = np.identity(3, dtype=np.int64)[y]
X_train, X_test, y_train, y_test = train_test_split(data.values, y, train_size=0.67, shuffle=True)
print(f"X_train: {X_train.shape}, X_test: {X_test.shape}, y_train: {y_train.shape}, y_test: {y_test.shape}")
X_train = torch.from_numpy(X_train).float()
y_train = torch.from_numpy(y_train).float()
X_test = torch.from_numpy(X_test).float()
y_test = torch.from_numpy(y_test).float()


# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 25, X_train.shape[1], 3, y_train.shape[1]

まずは適当に学習させてみる

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(TwoLayerNet, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        In the forward function we accept a Tensor of input data and we must return
        a Tensor of output data. We can use Modules defined in the constructor as
        well as arbitrary operators on Tensors.
        """
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

# Construct our model by instantiating the class defined above
model = TwoLayerNet(D_in, H, D_out)

# Construct our loss function and an Optimizer. The call to model.parameters()
# in the SGD constructor will contain the learnable parameters of the two
# nn.Linear modules which are members of the model.
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)

losses = []
testlosses = []

for t in range(50):
    # Forward pass: Compute predicted y by passing x to the model

    y_pred = model(X_train)

    # Compute and print loss
    loss = criterion(y_pred, y_train)
    losses.append(loss.item())
    if t % 10 == 0:
        print(t, loss.item())

    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    with torch.no_grad():
        y_test_pred = model(X_test)
        testloss = criterion(y_test_pred, y_test)
        testlosses.append(testloss)
plt.plot(losses)
plt.plot(testlosses)

無事学習は進んでいるようです。
モデルもexampleから適当に見繕いました。
※分類問題なのでクロスエントロピーを利用するのがベストプラクティスかもしれませんが、なぜかうまくいかなかったことと本題から逸れるため、このままMSEで行きます。

forwardを理解したい

モデルの中のforwardを理解したい

上のコードの中から抜粋

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(TwoLayerNet, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        In the forward function we accept a Tensor of input data and we must return
        a Tensor of output data. We can use Modules defined in the constructor as
        well as arbitrary operators on Tensors.
        """
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

`torch.nn.Linear` は重みとバイアスだけを持つ簡単な結合層です。(Linear — PyTorch 1.6.0 documentation)

隠れ層は一層だけで、今回適当に3を設定しています。入力層・隠れ層・出力層をこの2行で表現しています。

        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

次に本題のforwardです。

    def forward(self, x):
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

forwardはデータxをtensor型で受け取って、3層のレイヤーを通って出力しているだけです。
clampはデータに最小値・最大値を与えてその範囲内にデータを納めます。（torch.clamp — PyTorch 1.6.0 documentation）
以下がexampleです。

>>> a = torch.randn(4)
>>> a
tensor([-1.7120,  0.1734, -0.0478, -0.0922])
>>> torch.clamp(a, min=-0.5, max=0.5)
tensor([-0.5000,  0.1734, -0.0478, -0.0922])

上記の学習コードでは最低値を0として他をそのままの値で出力しています。つまりReLU関数と同等です。そのための出力の変数もh_reluとなっています。

なぜforwardに書く必要があるのか？

本題です。私はこれに悩みました。
特にこの後実験しようとしている損失関数や活性化関数の中でもforwardメソッドは定義されます。
この辺りがすごく不思議に思えてしまったのです。

ここでニューラルネットワークの基本的な構造を思い出します。

ニューラルネットワークの基礎を学んだなら当たり前のことなのですが、それぞれの層・活性化関数・損失関数も全て入力から出力までを順伝搬する仕組みです。
全てデータを左から受け取って右に流していく役割を持っているので、定義するのは至極当たり前でした。
それぞれをレイヤーとして認識してnn.Moduleクラスを基底として全てに順伝搬を定義しているだけです。

さらに深く理解するため、backwardの役割を改めて見ておきます。

optimizer.zero_grad()
loss.backward()
optimizer.step()

ニューラルネットワークの重みを更新するためには誤差逆伝搬法を利用します。（誤差逆伝搬法が何かについては探せばいくらでも記事が出てきますのでこちらでは割愛）
端的にいうとそれぞれの重みの勾配を計算しそれが損失関数にどれくらい影響を与えているのか計算できます。
重みの更新には多くの手法があり上記ではSGDを利用しています。

それでは上のコードが何をしているか確認していきます。
optimizer.zero_grad()は一旦「計算した勾配の初期化」だと思っておいて下さい。

次にloss.backwordを見ていきます。
lossは損失関数でここではMSE(Mean Square Error)を利用しています。ですが、損失関数と名のついた関数には変わりありません。
別の簡素な関数を用意して、事例を見ていきます。

backwardは何をしているのか。PytochのAutogradという概念。

x = torch.tensor(3.0, requires_grad=True)

簡単な関数を用意しました。 x = 3です。これを入力だと意識します。
requires_gradは勾配を自動で計算することを定義する引数です。ここでTrueとしておくと、その先にある様々の層の計算に大して、どれくらい寄与するのかその勾配を計算します。

そして次に出力としてy = 2xを定義します。

y = 2*x
print(y)
# tensor(4., grad_fn=<MulBackward0>)

この時y = 2xのxの勾配を出力します。
その値はx.gradに入っています。ただし先に勾配を計算する必要があります。

y.backward()

です。こうすることでx.gradは計算されます。xを定義する時にrequires_grad=Trueとしていないと、このタイミングでエラーが発生します。
結果はy = 2xにおけるxの勾配は2となります。

print(x.grad)
# tensor(3.)

次に伝搬のイメージをするために新しくz=5yを用意します。

z = 5*y

これは合成関数の微分を使って解くと（代入してもさくっと解けてしまいますが）zをxで微分すると15となります。これがzに対するxの勾配です。
あたらためて全て書き直すと以下のようになります。

x = torch.tensor(2.0, requires_grad=True)
y = 3*x
z = 5*y
z.backward()
print(x.grad)
# tensor(15.)

なぜ全て書き直したかというと、この時４行目のz.backwardをy.backwardにすると3が出力されるからです。ぜひ丸っとコピペして試してみてください。
つまりbackwardはrequires_grad=Trueとした変数に対して（ここではx）、目的の関数（ここではy, z）に対して微分を行った時の勾配を計算する事ができます。

（個人的に、まだ自分自身が慣れてないだけですが、z.gradとかy.gradと表示された方がみやすい気がしてしまいます。ただ実際のニューラルネットワークの各層に対してこの勾配を持つので、各層の変数から取り出せた方がいいというのも理解できます。）

そしてこの更新に対して注意が必要なのは、計算された重みは追加されていくという事です。つまりz.backward()をn実行すると、x.gradはn * 15となります。
そのためoptimizer.zero_grad()で勾配を初期化する必要が出てくるわけです。

最後にoptimizer.step()

これは重みの更新を行っています。
この行をコメントアウトすると学習の結果は以下のようになります。

何も更新されません。

optimizer.step()で学習率と最適化手法に基づいて重みを更新しているわけです。
さらにschedulerなどを設定している時は、学習率の動的な変更もこのタイミングで行われています。

損失関数を新しく定義したい時もforward

class RMSELoss(nn.Module):
    def __init__(self, eps=1e-6):
        super().__init__()
        self.mse = nn.MSELoss()
        self.eps = eps

    def forward(self, yhat, y):
        loss = torch.sqrt(self.mse(yhat, y) + self.eps)
        return loss

これはMSEをRMSEに変更したい時のコードです。
（ちなみにこれを見て理解はできるものの本当に説明できるか不安になり、本記事につながりました）
これを利用し、上記のコードで学習し直すと。

無事学習できました！
RMSEはMSEの平方根を取る形です。MSEに比べて外れ値に鈍感になる上損失は数字上小さくなるので（単位も異なりますので）学習を回す回数は増えましたが、ちゃんと収束しました。

活性化関数を新しく定義したい時もforward

class ReLU(Module):
    def __init__(self, inplace: bool = False):
        super(ReLU, self).__init__()
        self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
        return F.relu(input, inplace=self.inplace)

まとめ

• PyTorchはnn.Moduleクラスを基底とし、順伝搬の処理をforwardの中に書いている。
• さらにnn.Moduleを基底として、それらの入力層・隠れ層・出力層・活性化関数・損失関数などを組み合わせる（つなぎ合わせる）ことでモデルを作成し学習する。
• 学習はAutogradという勾配を自動で計算する仕組みに支えられている。
• そのAutogradを利用して、optimizer.zero_grad()で勾配を初期化し、loss.barkward()でそのエポックにおける損失に対する勾配を計算し、optimizer.step()で重みの更新を行う。この繰り返しで学習が進む。

ご覧いただきありがとうございました！

下記書籍とブログも参考にしています。
つくりながら学ぶ！PyTorchによる発展ディープラーニング | 小川 雄太郎 | 工学 | Kindleストア | Amazon
PyTorchは誤差逆伝播とパラメータ更新をどうやって行っているのか？ - け日記