【PyTorchチュートリアル】LEARNING PYTORCH WITH EXAMPLESの6、nnモジュールでクラス化

PyTorchのチュートリアルの「LEARNING PYTORCH WITH EXAMPLES」を進めていきます。今回は、「Cumtom nn Modules」から行なっていきます。

これで、いつもチュートリアルで見ることができる様なネットワークのプログラム構造になりました。

こんな人の役に立つかも

・機械学習プログラミングを勉強している人

・PyTorchで機械学習プログラミングを勉強している人

・PyTorchでニューラルネットワークを勉強している人

nnモジュールのクラス化

まずは、torchのimportです。

# -*- coding: utf-8 -*-
import torch

次に、クラスで２層ニューラルネットワークを定義します。今回作成するクラスは、クラス名「TwoLayerNet」で、torch.nn.Moduleを継承して作成します。

class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        コンストラクタ
        """
        super(TwoLayerNet, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        順伝播の定義。xが入力するテンソル。
        """
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

コンストラクタ（初期化関数）で「linear1」と「linear2」を定義しています。

メソッドとして、「forward」を定義します。これは順伝播に利用するメソッドになります。

プログラムの流れとしては、linear1に入力して、その出力を活性化関数ReLUで変換（実際にはclump(min=0)）したものをh_relu変数に格納します。そして、h_reluをlinear2に入力します。その出力値がy_pred変数に格納されて、forwardメソッドからの返り値になります。

次に、定義したクラスを利用してニューラルネットワークの訓練を行います。

# 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)

# 誤差関数：平均二乗誤差
criterion = torch.nn.MSELoss(reduction='sum')
# 最適化手法：勾配降下法：学習率lrもここで定義
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for t in range(500):
    # ①順伝播
    y_pred = model(x)

    # ②誤差の計算
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # ③逆伝播をして、重みの更新を行う。
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

今回、nnモジュールを継承したクラスでネットワークを作成することで、①順伝播の部分に直接順伝播の計算式を記述する必要がなくなりました。

順伝播はクラスのところで記載される様になったので、より抽象的にプログラムを記載できる様になりました。

制御と重みの共有について

この項目は、理解するのに少し時間がかかりました。いきなり奇妙なネットワークを実装する、と始まりますので・・・

この項目では、ネットワークのモデルの順伝播を定義するときに、Pythonのループ構造や制御文を利用することができますよ、という内容だと理解できます。

また、ループでモジュール（チュートリアルではmiddle_linear）を使い回しても、計算グラフはしっかりとそれぞれで計算されるので大丈夫ですよ、という内容だと理解しました。

# -*- coding: utf-8 -*-
import random
import torch

class DynamicNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        コンストラクタ：「middle_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):
        """
        順伝播の「middle_layer」を0〜3回forループで繰り返します。
        「middle_layer」をループ回数使い回しています。
        
        順伝播を定義するときには、Pythonのループや制御文を利用することができます。
        
        ここでは、計算グラフを定義するときに同じモジュールを何度も再利用しても大丈夫なことを示しています。
        これはLuaTorch(昔のPyTorch)からの改善点で、昔は各モジュールは１回しか使えませんでした。
        """
        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)
        y_pred = self.output_linear(h_relu)
        return y_pred

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)

criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
for t in range(500):
    # 順伝播
    y_pred = model(x)

    # 損失の計算
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # 最適化と重みの更新
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

今回のプログラムをGitHubにアップロードしていますので、ご利用ください。

チュートリアルの記事

以下に、「LEARNING PYTORCH WITH EXAMPLES」チュートリアルの一連の記事をリンクします。