🧠

[PyTorch] ニューラルネットワークをゼロから実装だけど勾配計算はライブラリに頼って楽をする

2022/09/09に公開約11,600字1件のコメント

前回:https://zenn.dev/mory22k/articles/e3a51a259e590e

前回の記事では、NumPy を使用して基本的にゼロからニューラルネットワークを実装しました。

https://zenn.dev/mory22k/articles/e3a51a259e590e

最大の問題点は、レイヤーごとに誤差逆伝播の計算式を手計算で導出する必要があることです。例えば全結合層では、次のように backward() を実装していました (若干変更を加えています):

class Affine:
    def __init__(self, input_size, output_size, init_std = 1):
        self.params = {
            'w': init_std * np.random.randn(input_size, output_size),
            'b': init_std * np.random.randn(output_size)
        }
        self.grads = {
            'w': None,
            'b': None
        }
        self.x = None

    def forward(self, x):
        if x.ndim==1:
            self.x = x.reshape(1, x.size)
        else:
            self.x = x
        y = self.x @ self.params['w'] + self.params['b']
        return y

    def backward(self, dy):
        dx = dy @ self.params['w'].T
        self.grads['w'] = self.x.T @ dy
        self.grads['b'] = np.sum(dy, axis=0)
        return dx

これでも問題なく学習が進むには進むのですが、レイヤーを作成するたびに逆伝播を定義するのが面倒です。各層の計算式を Y = f(X) の形に書いた上で、更に L = g(Y) を仮定し、それを X で微分する必要があるのです:

L = g(Y) = g \circ f(X)
\Downarrow
\frac{\partial L}{\partial X} = \frac{\partial Y}{\partial X} \frac{\partial L}{\partial Y}

PyTorch の自動微分機能を用いれば、この過程を完全にすっ飛ばすことができます。基本的には順伝播の設計のみに注力するだけで済むようになります。

そういうわけで、この記事では前回作成した浅い NN を PyTorch で実装します。

import torch
from matplotlib import pyplot as plt

自動微分の基本

PyTorch の自動微分は 死ぬほど 便利ですが、使い方にちょっとだけ工夫が要ります。まずはシンプルに

y = \sin(x)

の微分計算を行ってみます。

x = torch.tensor(torch.pi / 2, dtype=float, requires_grad = True)
y = torch.sin(x)

# 自動微分
y.backward()

# 結果表示
print(int(y.item()))
print(int(x.grad.item()))
1
0

解析計算してみると

\begin{align*} y(x) &= \sin(x) & \rightarrow && y\!\left(\frac{\pi}{2}\right) &= 1 \\ \frac{dy}{dx} &= \cos(x) & \rightarrow && \left. \frac{dy}{dx} \right|_{x = \pi / 2} &= 0 \end{align*}

なので、たしかに正しい結果になっています。

グラフも描いてみます。しかし、安直に次のようにやろうとすると、

x = torch.linspace(0, 2*torch.pi, 500, requires_grad=True)
y = torch.sin(x)
y.backward()
RuntimeError: grad can be implicitly created only for scalar outputs

怒られます。

エラーメッセージの通り、backward() はスカラー値にしか実行できません。なので、次のようにします。

x = torch.linspace(0, 2*torch.pi, 500, requires_grad=True)

y = []
for x_tmp in x:
    y_tmp = torch.sin(x_tmp)
    y.append(y_tmp.item())
    y_tmp.backward()
xx = x.detach()

plt.plot(xx, y, label='y')
plt.plot(xx, x.grad, label='dy/dx')
plt.legend()
plt.grid()
plt.xlabel('x')
plt.show()

複雑な関数に対しても勾配計算ができます。

def func(x):
    y = (2*x**2 + 4*torch.sin(2 * x))**torch.exp(-x/5)
    return y

x = torch.linspace(0, 2*torch.pi, 500, requires_grad=True)

y = []
for x_tmp in x:
    y_tmp = func(x_tmp)
    y.append(y_tmp.item())
    y_tmp.backward()

xx = x.detach()

plt.plot(xx, y, label='y')
plt.plot(xx, x.grad, label='dy/dx')
plt.legend()
plt.grid()
plt.xlabel('x')
plt.show()

これをフル活用して勾配計算を楽に行おうというのが今回の趣旨です。

レイヤーの定義

全結合層

Y = X \cdot W + b

前回と同様に、X(n, d_1) 型配列、W(d_1, d_2) 型配列、b(d_2) 型配列、Y(n, d_2) 型配列を想定しています。

\begin{align*} & X: && (n, d_1) \\ & W: && (d_1, d_2) \\ & b: && (d_2) \\ & Y: && (n, d_2) \end{align*}
class Affine:
    def __init__(self, input_size, output_size, init_std = 1):
        # パラメータ定義
        self.w = init_std * torch.randn(input_size, output_size)
        self.b = init_std * torch.randn(output_size)

        # 自動微分有効化
        self.w = self.w.requires_grad_(True)
        self.b = self.b.requires_grad_(True)

        # リストにまとめておく
        self.params = [self.w, self.b]

    def forward(self, x):
        y = x @ self.w + self.b
        return y

各パラメータ (w, b) の定義が完了した後に .requires_grad_(True) によって逆伝播計算を有効化しています。

実際に計算してみます。

# レイヤー定義
fc1 = Affine(2, 4)
fc2 = Affine(4, 1)

# 順伝播
x = torch.tensor([[1., 2.]], requires_grad=True) # shape (1, 2)
x = fc1.forward(x)
x = fc2.forward(x)

# 逆伝播
x.backward()

# 結果
print(f'=== RESULT ===')
print(x.item())
print()
print(f'=== FC 1 ===')
print(fc1.w.grad)
print(fc1.b.grad)
print()
print('=== FC 2 ===')
print(fc2.w.grad)
print(fc2.b.grad)
=== RESULT ===
2.487516403198242

=== FC 1 ===
tensor([[ 0.1626,  0.2174,  0.3500, -0.4937],
        [ 0.3252,  0.4349,  0.7001, -0.9873]])
tensor([ 0.1626,  0.2174,  0.3500, -0.4937])

=== FC 2 ===
tensor([[-3.6132],
        [ 2.6920],
        [ 6.9292],
        [-1.8240]])
tensor([1.])

ReLU 層

f_\mathrm{ReLU}(x) = \max(x, 0)
class ReLU:
    def __init__(self):
        pass

    def forward(self, x):
        y = torch.maximum(x, torch.tensor(0.))
        return y

先程の計算例の fc1fc2 の間に ReLU を挟んでみます。

# レイヤー定義
fc1 = Affine(2, 4)
fc2 = Affine(4, 1)
relu = ReLU()

# 順伝播
x = torch.tensor([[1., 2.]], requires_grad=True) # shape (1, 2)
x = fc1.forward(x)
x = relu.forward(x)
x = fc2.forward(x)

# 逆伝播
x.backward()

# 結果
print(f'=== RESULT ===')
print(x.item())
print()
print(f'=== FC 1 ===')
print(fc1.w.grad)
print(fc1.b.grad)
print()
print('=== FC 2 ===')
print(fc2.w.grad)
print(fc2.b.grad)
=== RESULT ===
-0.9565512537956238

=== FC 1 ===
tensor([[ 0.0000, -0.1650,  0.2169, -0.1234],
        [ 0.0000, -0.3301,  0.4338, -0.2468]])
tensor([ 0.0000, -0.1650,  0.2169, -0.1234])

=== FC 2 ===
tensor([[0.0000],
        [0.1015],
        [1.2226],
        [2.1728]])
tensor([1.])

前回、頑張って手計算で逆伝播の計算式を求めていたのがばかばかしく感じられるほど楽です。

平均二乗誤差

\mathcal L_{\textrm{MSE}} = \frac{1}{2n} \sum_{i=1}^{n} (Y_i - T_i)^2
\begin{align*} \textrm{expected:} && Y &= \begin{bmatrix} Y_1 & Y_2 & \cdots & Y_n \end{bmatrix} \\ \textrm{true:} && T &= \begin{bmatrix} T_1 & T_2 & \cdots & T_n \end{bmatrix} \end{align*}

予測結果と実際の値の差の二乗の和をデータ数 n で割ったものです。

class MSE:
    def __init__(self):
        pass

    def forward(self, pred, true):
        n = pred.shape[0]
        loss = torch.sum((pred-true)**2) / n / 2
        return loss

適当に教師データを設定して計算してみましょう。

# レイヤー定義
fc1 = Affine(2, 4)
fc2 = Affine(4, 1)
relu = ReLU()
loss_func = MSE()

# 入力と教師データ
x = torch.tensor([[1., 2.]])
t = torch.tensor([[4.]])

# 順伝播
x = x.clone().requires_grad_(True)

x = fc1.forward(x)
x = relu.forward(x)
x = fc2.forward(x)

# 損失を計算
loss = loss_func.forward(x, t)

# 逆伝播
loss.backward()

# 結果
print(f'=== RESULT ===')
print(x.item())
print()
print(f'=== LOSS ===')
print(loss.item())
print()
print(f'=== FC 1 ===')
print(fc1.w.grad)
print(fc1.b.grad)
print()
print('=== FC 2 ===')
print(fc2.w.grad)
print(fc2.b.grad)
=== RESULT ===
-2.519716262817383

=== LOSS ===
21.25334930419922

=== FC 1 ===
tensor([[ 5.5313,  0.0000,  0.0000,  0.0000],
        [11.0626,  0.0000,  0.0000,  0.0000]])
tensor([5.5313, 0.0000, 0.0000, 0.0000])

=== FC 2 ===
tensor([[-13.6530],
        [  0.0000],
        [  0.0000],
        [  0.0000]])
tensor([-6.5197])

勾配法による回帰

教師データ

まずは教師データを準備しましょう。

x_train = torch.rand(500) * 16 - 8
x_train = x_train.reshape(len(x_train), 1)
y_train = x_train**3

ネットワークの設計

続いてネットワークを設計します。

class SimpleNet:
    def __init__(self):
        self.fc1 = Affine(1, 64)
        self.fc2 = Affine(64, 1)
        
        self.relu = ReLU()

        self.params = self.fc1.params + self.fc2.params
    
    def forward(self, x):
        x = self.fc1.forward(x)
        x = self.relu.forward(x)
        x = self.fc2.forward(x)
        return x

Network = SimpleNet()

学習方法の設定

学習メソッドと学習のためのアルゴリズムを設計します。

損失には MSE を、学習には Adam を使用します。今回は PyTorch に準備されている Adam オプティマイザを利用します。

loss_func = MSE()
optimizer = torch.optim.Adam(
    params= Network.params,
    lr=     1e-3
)

学習のためのアルゴリズムは次のとおりです。

1. epoch数だけ繰り返し:
2.     順伝播を計算
3.     損失を計算
4.     勾配を計算
5.     1ステップ学習
6.     損失を記録

これを PyTorch で書き表します。

def optimize(Network, loss_func, optimizer, x_train, y_train, epochs):
    true = y_train
    loss_hist = []
    for epoch in range(epochs):
        # 順伝播を計算
        pred = Network.forward(x_train)

        # 損失を計算
        loss = loss_func.forward(pred, true)

        # 勾配を計算
        optimizer.zero_grad()
        loss.backward()

        # 1ステップ学習
        optimizer.step()

        # 損失を記録
        loss_hist.append(loss.detach().item())
        print(f'epoch: {epoch}, loss: {loss.detach().item()}')

    return loss_hist
loss_hist = optimize(Network, loss_func, optimizer, x_train, y_train, epochs=8000)

学習結果を見てみましょう。

plt.plot(loss_hist)
plt.show()

with torch.no_grad():
    x = torch.linspace(-10, 10, 1000).reshape(1000, 1)
    y = Network.forward(x)

plt.plot(x, y)
plt.scatter(x_train, y_train, s=10, color='#000000', marker='x')
plt.show()

ミニバッチ学習

続いてミニバッチによる学習もやってみます。

教師データ分割

まず教師データをいくつかのデータに分割します。

batch_size = 25
batch_num = len(x_train) // batch_size
for i in range(batch_num):
    print(f'batch{ i }:\n{ batch_size*i } ~ { batch_size*(i+1) }')
batch0:
0 ~ 25
batch1:
25 ~ 50
...
batch19:
475 ~ 500

これらのデータをすべて参照することを1エポックと呼びます。今回は20個のバッチに分割したので、1エポックで20回学習が行われることになります。

学習方法の設定

def optimize_batch(Network, loss_func, optimizer, x_train, y_train, epochs, batch_size=None):
    if batch_size is None:
        batch_size = len(x_train)
    batch_num = len(x_train) // batch_size
    
    loss_hist = []
    for epoch in range(epochs):
        loss_hist_hist = [] 
        for i in range(batch_num):
            x_train_batch = x_train[batch_size * i : batch_size * (i+1)]
            y_train_batch = y_train[batch_size * i : batch_size * (i+1)]

            pred = Network.forward(x_train_batch)

            loss = loss_func.forward(pred, y_train_batch)

	    optimizer.zero_grad()
            loss.backward()

            optimizer.step()

            loss_hist_hist.append(loss.detach().item())

        # 損失を記録
        loss_hist_hist_mean = torch.mean(torch.tensor(loss_hist_hist))
        loss_hist.append(loss_hist_hist_mean)
        print(f'epoch: {epoch}, loss: {loss_hist_hist_mean}')

    return loss_hist
Network = SimpleNet()
loss_func = MSE()
optimizer = torch.optim.Adam(
    params= Network.params,
    lr=     1e-4
)

loss_hist = optimize_batch(
                Network,
                loss_func,
                optimizer,
                x_train,
                y_train,
                epochs = 2000,
                batch_size = 25
            )

plt.plot(loss_hist)
plt.show()

with torch.no_grad():
    x = torch.linspace(-10, 10, 1000).reshape(1000, 1)
    y = Network.forward(x)

plt.plot(x, y)
plt.scatter(x_train, y_train, s=10, color='#000000', marker='x')
plt.show()

まとめ

PyTorch の自動微分めっちゃ便利。癖になりそう。

参考文献

  1. 斎藤 康毅, “ゼロから作るDeep Learning ―Pythonで学ぶディープラーニングの理論と実装,” オライリージャパン, 2016.
  2. 岡谷 貫之, “深層学習 改訂第2版,” 機械学習プロフェッショナルシリーズ, 講談社, 2022.

Discussion

2022/9/11 修正

loss.backward() を実行する前に optimizer.zero_grad() を実行するのを忘れるという、マジで絶対にやってはいけないミス を堂々とやっていたため色々とバグっていました。
なのでそこを修正しました。

ログインするとコメントできます