【PyTorchチュートリアル⑤】Learning PyTorch with Examples （後編）

はじめに

前回に引き続き、PyTorch 公式チュートリアル の第5弾です。
今回は Learning PyTorch with Examples の後編です。
前編はこちらです。

3. nn module

3.1. PyTorch: nn

autograd だけでは、ニューラルネットワークのモデルを作成することはできません。
モデルの構築は、nnパッケージを利用します。
nnパッケージには、入力層、隠れ層、出力層を定義する Sequential クラスや、損失関数も含まれています。
下の例では、nnパッケージを使用して2層ネットワークを実装します。

import torch

# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# ランダムな入力データと教師データを作成します
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

# Use the nn package to define our model as a sequence of layers. nn.Sequential
# is a Module which contains other Modules, and applies them in sequence to
# produce its output. Each Linear Module computes output from input using a
# linear function, and holds internal Tensors for its weight and bias.

# nnパッケージを使用して、モデルを一連のレイヤーとして定義します。 
# nn.Sequentialは、他のモジュールを順番に適用してモデルを生成します。
# Linear は、線形関数を使用して入力から出力を計算し、重みとバイアスを
# 内部テンソルで保持します。
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out),
)

# nnパッケージには、損失関数も含まれています。
# 今回は、損失関数として平均二乗誤差（MSE）を使用します。
loss_fn = torch.nn.MSELoss(reduction='sum')

learning_rate = 1e-4
for t in range(500):
    # 順伝播：モデルにxを渡すことにより、予測値 y を計算します。
    # 基底クラスである Module は__call__演算子をオーバーライドするため、
    # 関数のように呼び出すことができます。
    # その場合、入力データのテンソルをモジュールに渡し、出力データのテンソルを生成します。
    y_pred = model(x)

    # ロス（損失）を計算し出力します。
    # 損失関数に y の予測値と教師データのテンソルを渡すと
    # 損失を含むテンソルを返します。
    loss = loss_fn(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # 逆伝播を計算する前に、勾配をゼロにします。
    model.zero_grad()

    # 逆伝播：損失の勾配を計算します。
    # 内部的には、モジュールのパラメーターはrequire_grad = True で保持されているため
    # すべてのパラメーターの勾配を計算します。
    loss.backward()

    # 確率的勾配降下法を使用して重みを更新します。

    with torch.no_grad():
        for param in model.parameters():
            param -= learning_rate * param.grad

3.2. PyTorch: optim

ここまでは、モデルの重みを以下のように自分で計算して更新していました。

param -= learning_rate * param.grad

この計算方法は確率的勾配降下法と呼ばれます。
モデルの重みの計算方法（ optimizer / 最適化アルゴリズム）は他にもあり、AdaGrad、RMSProp、Adam など、より高度なアルゴリズムを使用したいケースもでてきます。
PyTorchのoptimパッケージには、いろいろな最適化アルゴリズムがあります。
以下の例では、上の例のように nn パッケージでモデルを定義しますが、重みの更新に optimパッケージのAdamアルゴリズムを使用します。

import torch

# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# ランダムな入力データと教師データを作成します
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

# nnパッケージを使用して、モデルと損失関数を定義します。
model = torch.nn.Sequential(
    torch.nn.Linear(D_in, H),
    torch.nn.ReLU(),
    torch.nn.Linear(H, D_out),
)
loss_fn = torch.nn.MSELoss(reduction='sum')

# optimパッケージを使用して、モデルの重みを更新する最適化アルゴリズム（オプティマイザ）を定義します。
# ここでは Adam を使用します。
# optim パッケージには、他にもいろいろな最適化アルゴリズムが含まれています。 
# Adam コンストラクタへの最初の引数は、どの Tensor を更新するかを指定します。
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for t in range(500):
    # 順伝播：モデルにxを渡すことにより、予測値 y を計算します。
    y_pred = model(x)

    # ロス（損失）を計算し出力します。
    loss = loss_fn(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # 逆伝播の計算前に、更新する変数（重み）のすべての勾配をゼロにします。
    # これは、backward が呼び出されるたびに、勾配が累積される（上書きされない）ためです。
    optimizer.zero_grad()

    # 逆伝播：損失の勾配を計算します。
    loss.backward()

    # オプティマイザのステップ関数を呼び出すと、パラメーターが更新されます。
    optimizer.step()

3.3. PyTorch: Custom nn Modules

複雑なモデルを構築したい場合、nn.Moduleをサブクラス化して実現できます。
サブクラスで forward 関数をオーバーライドし、入力 Tensor から出力 Tensor を返却する処理を記述することで独自のモジュールを定義できます。
この例では、2層ネットワークをカスタムモジュールサブクラスとして実装します。

import torch

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        コンストラクタで、2つの nn.Linear モジュールをインスタンス化し、
        それらをメンバー変数として割り当てます。
        """
        super(TwoLayerNet, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        forward 関数では、入力 Tensor を元に、出力 Tensor を返す必要があります。
        コンストラクタで定義されたモジュールを使用できます。
        """
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred


# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# ランダムな入力データと教師データを作成します
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

# 上で定義したニューラルネットワークモジュールをインスタンス化してモデルを構築します
model = TwoLayerNet(D_in, H, D_out)

# 損失関数と最適化アルゴリズム(オプティマイザ)を定義します。 
# SGD の引数 model.parameters() には、定義したクラスのメンバーである
# 2つの nn.Linear モジュールのパラメーターも含まれます。
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for t in range(500):
    # 順伝播：モデルに x を渡すことにより、予測値 y を計算します。
    y_pred = model(x)

    # ロス（損失）を計算し出力します。
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # 勾配をゼロにし、逆伝播を計算して、重みを更新します。
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

3.4. PyTorch: Control Flow + Weight Sharing

動的グラフと重み共有(Weight Sharing)の例として、特殊なモデルを実装します。以下のReLUネットワークは、forward 関数で0から3の乱数を選択し、複数の隠れ層で同じ重みを共有して計算します。
このモデルをModuleサブクラスとして実装します。

import random
import torch


class DynamicNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        コンストラクターでは、forward 関数で使用する3つの nn.Linear インスタンス(入力層、隠れ層、出力層)を作成します。
        """
        super(DynamicNet, self).__init__()
        self.input_linear = torch.nn.Linear(D_in, H)
        self.middle_linear = torch.nn.Linear(H, H)
        self.output_linear = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        forward 関数では、0回 から 3回 の値をランダムに選択し、
        middle_linearモジュールを複数回再利用して、隠れ層の処理を計算します。

        autograd は動的計算グラフのため、順伝播時に構築されます。
        そのため、forward 関数内では、ループや条件ステートメントなど、通常のPythonの処理を記述できます。

        計算グラフを定義するときに同じモジュールを何度も利用することができます。
        これはLua Torchからの改善点で、Lua Torch では各モジュールは1回しか使用できませんでした。
        """
        h_relu = self.input_linear(x).clamp(min=0)
        for _ in range(random.randint(0, 3)):
            h_relu = self.middle_linear(h_relu).clamp(min=0)
            print(str(_))
        print(h_relu.size())
        y_pred = self.output_linear(h_relu)
        return y_pred


# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# ランダムな入力データと教師データを作成します
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

# 上で定義したニューラルネットワークモジュールをインスタンス化してモデルを構築します
model = DynamicNet(D_in, H, D_out)

# 損失関数とオプティマイザーを作成します。
# このモデルは、通常の確率勾配降下法で学習させる(収束させる)のは難しいので、momentumを指定します。
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
for t in range(500):
    # 順伝播：モデルにxを渡すことにより、予測値 y を計算します。
    y_pred = model(x)

    # ロス（損失）を計算し出力します。
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # 勾配をゼロにし、逆伝播を計算して、重みを更新します。
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

終わりに

以上が、PyTorch の5つ目のチュートリアル「Learning PyTorch with Examples」の内容です。
autograd、torch.nn パッケージ、torch.optim パッケージの理解を深めることができました。

次回は6つ目のチュートリアル「What is torch.nn really?」を進めてみたいと思います。

履歴

2020/07/10 初版公開
2020/07/10 前編のリンク修正
2020/10/10 次回のリンク追加