PyTorchでネットワークを実装する

引き続きPyTorchのお勉強中です。

前々回はテンソル、前回は誤差逆伝播について見ていきましたが、今回はtorch.nnのモジュールを中心にネットワークの作り方について整理していきます。

前回より少し難しくして、2次関数 (w=4.5, b=7.0, εは誤差) を例にしていきます。

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

print(torch.__version__)  # 1.1.0

np.random.seed(42)

w2 = 0.6
w1 = 4.5
b = 7.0

x_array = np.array([-1.5, -1.0, -0.1, 0.9, 1.8, 2.2, 3.1])
y_array = w2 * x_array**2 + w * x_array + b + np.random.normal(size=x_array.shape[0])
plt.scatter(x_array, y_array)

x = torch.tensor(x_array).unsqueeze(1).float()
print(x)  # tensor([[-1.5000], [-1.0000], [-0.1000], [0.9000], [1.8000], [2.2000], [3.1000]])
y = torch.tensor(y_array).unsqueeze(1).float()
print(y)  # tensor([[2.0967], [2.9617], [7.2037], [13.0590], [16.8098], [19.5699], [28.2952]])

モジュールとtorch.nn

PyTorchにおいて、畳み込みフィルタ・活性化関数・損失関数などのネットワークの"層"を実装しているのが、モジュールです。このモジュールを組み合わせることによって、ネットワークを実装します。

PyTorchではよく使われる層はモジュールとして既に実装されていることが多く、それらはtorch.nnに集約されています。

モジュールの基本的な使い方

最初に の1次元関数でモデリング・学習するコードを実装し、モジュールの基本的な使い方を確認します。

以下がそのコードです。2つのモジュールを使っており、SGDで学習してます。学習のループの詳細は前回の投稿も参考にしてみてください。

• nn.Linearは1次関数で、傾きと切片がパラメータとなります
  • __init__引数は入力サイズ・出力サイズ
• nn.MSELossは平均2乗誤差を計算するモジュールで、損失関数として利用しています

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=1e-3)

for epoch in range(1, 2001):
    y_p = model(x)
    loss = loss_fn(y_p, y)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch == 1 or epoch % 200 == 0:
        print(f"Epoch {epoch}: loss={loss}, params={list(model.parameters())}")

# Epoch 1: loss=247.5394287109375, params=[Parameter containing:tensor([[-0.3880]], requires_grad=True), Parameter containing:tensor([0.6815], requires_grad=True)]
# ...
# Epoch 2000: loss=2.2267117500305176, params=[Parameter containing:tensor([[5.5755]], requires_grad=True), Parameter containing:tensor([8.3513], requires_grad=True)]

概ね前回の投稿と同じですが、torch.nnのモジュールの特徴を2点補足します。

parametersメソッド

学習可能なパラメータはparametersメソッドで取得できます。

optim.SGDなどのOptimizerを継承した最適化クラスをインスタンス化する際に、第1引数 (params) にparametersメソッドから得られた値を渡すことで勾配計算 (学習) の対象に含めることができます。

print(list(model.parameters()))
# [Parameter containing: tensor([[5.5755]], requires_grad=True),
#  Parameter containing: tensor([8.3513], requires_grad=True)]

モジュールのパラメータはParameterオブジェクトとしてプロパティに持っています。実体はtensorなのですが、parametersメソッドはこのクラスのプロパティのみを探索しています。

• Linearでは2つのプロパティ (weightとbias) がParameterオブジェクトです (参考)

0次元目はバッチのインデックス

nn.Moduleは0次元目がバッチ内のインデックスが入っていることを前提としてます。そのため以下のケースはエラーとなります。モデルは入力サイズ1を期待しているのに、実際の入力サイズは7のためです。
上のコードではunsqueezeメソッドで1次元目 (サイズ1) を追加することで回避してます。

x = torch.tensor(x_array).float()  # unsqueeze無し
print(x.shape)  # torch.Size([7])

y_pred = model(x)
# ---------------------------------------------------------------------------
# RuntimeError                              Traceback (most recent call last)
# ...
# RuntimeError: size mismatch, m1: [1 x 7], m2: [1 x 1] at /opt/conda/conda-bld/pytorch_1556653114079/work/aten/src/TH/generic/THTensorMath.cpp:961

ネットワークの自作

上でLinearを使ってモデリングしましたが、学習の結果lossは下がりきりませんでした。もうちょっと複雑なモデルを導入する必要がありそうです。

そこで自前のネットワークを定義していくのですが、上の例で見たシンプルなインタフェース (__call__で順伝播、backwardで逆伝播) は維持したいところです。

LinearやMSELossなどのモジュール (TensorFlowではレイヤと呼ばれる) はnn.Moduleクラスを継承してます。

• https://pytorch.org/docs/stable/_modules/torch/nn/modules/linear.html#Linear
• https://pytorch.org/docs/stable/_modules/torch/nn/modules/loss.html#MSELoss

nn.Moduleクラスを継承したクラスを作成することで、シンプルなインタフェースで学習を実現できます。

以下コードでは、隠れ層 (1層) の簡単なネットワークをnn.Moduleを継承したMyModelで実装しています。使い方は上で見たLinearと全く同じことに着目してください。

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.linear_1 = nn.Linear(1, 5)
        self.linear_2 = nn.Linear(5, 1)
        
    def forward(self, x):
        x = self.linear_1(x)
        x = torch.relu(x)
        x = self.linear_2(x)
        return x

model = MyModel()
print(model)
# MyModel(
#   (linear_1): Linear(in_features=1, out_features=5, bias=True)
#   (linear_2): Linear(in_features=5, out_features=1, bias=True)
# )

loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=1e-3)

for epoch in range(1, 2001):
    y_p = model(x)
    loss = loss_fn(y_p, y)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch == 1 or epoch % 200 == 0:
        print(f"Epoch {epoch}: loss={loss}")
        
# Epoch 1: loss=238.2289276123047
# Epoch 200: loss=1.2844644784927368
# ...
# Epoch 2000: loss=1.1387771368026733
        
print(list(model.parameters()))
# [Parameter containing: tensor([[-0.9788], [0.6803], [1.7500], [1.0663], [0.5548]], requires_grad=True), 
#  Parameter containing: tensor([-0.5819, 0.9828, 0.7439, 1.3531, 1.1561], requires_grad=True), 
#  Parameter containing: tensor([[0.3001, 0.8646, 1.9179, 1.5838, 1.1189]], requires_grad=True), 
#  Parameter containing: tensor([1.4064], requires_grad=True)]

MyModelで実装しているのは__init__メソッドとforwardメソッドのみです。それぞれポイントを説明していきます。

__init__

__init__メソッドではネットワークを構成するパラメータとモジュールを初期化して、プロパティにセットしていきます。

モジュールのプロパティに対しても再帰的にパラメータを探索します。以下ではMyLinearsのインスタンスをプロパティに持つMyModelを定義してますが、print(model)でMyLinearの中のモジュールのパラメータまで認識されていることがわかります。

class MyLinears(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.linear1 = nn.Linear(1, 5)
        self.linear2 = nn.Linear(5, 1)
        
    def forward(self, x):
        x = self.linear1(x)
        x = torch.relu(x)
        x = self.linear2(x)
        return x

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.my_linears = MyLinears()
        self.linear = nn.Linear(1, 1)
        
    def forward(self, x):
        x = self.my_linears(x)
        x = torch.relu(x)
        x = self.linear(x)
        return x

model = MyModel()
print(model)
# MyModel(
#   (my_linears): MyLinears(
#     (linear1): Linear(in_features=1, out_features=5, bias=True)
#     (linear2): Linear(in_features=5, out_features=1, bias=True)
#   )
#   (linear): Linear(in_features=1, out_features=1, bias=True)
# )

ただしModuleを継承したプロパティのみがパラメータの探索対象となることに注意です。以下のようにlistなどでセットしてもパラメータ扱いされません。

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.linears = [nn.Linear(1, 5), nn.Linear(5, 1)]
        
    def forward(self, x):
        x = self.linears[0](x)
        x = torch.relu(x)
        x = self.linears[1](x)
        return x

model = MyModel()
print(model) # MyModel()
print(list(model.parameters()))  # []

ModuleListやModuleDictでプロパティにセットすると、パラメータを探索するようになります。

• それぞれ__getitem__や__len__などは実装されており、list・dictと同じアクセスができます

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.linears = nn.ModuleList([
            nn.Linear(1, 5), nn.Linear(5, 1)
        ])
        
    def forward(self, x):
        x = self.linears[0](x)
        x = torch.relu(x)
        x = self.linears[1](x)
        return x

model = MyModel()
print(model)
# MyModel(
#   (linears): ModuleList(
#     (0): Linear(in_features=1, out_features=5, bias=True)
#     (1): Linear(in_features=5, out_features=1, bias=True)
#   )
# )

forward

forwardメソッドで各モジュールを順につなげます。
上の実装ではxを介して 入力 -> Linear -> relu -> Linear -> 出力 という順番で受け渡ししてます。

ただしこのforwardメソッドを直接実行して順伝播させるのはNGで、__call__を使わないといけません。__call__内ではforward以外にhookも一緒に実行しているためです。

hookには順伝播時に実行したい任意の処理を定義できます。以下ではregister_forward_pre_hookとregister_forward_hookでprint関数をセットし、__call__で順伝播させることで各hookを実行させてます。

• register_forward_pre_hookメソッドで渡された関数は順伝播前、register_forward_hookメソッドで渡された関数は順伝播後に実行されます
• ちなみにregister_backward_hookもあります

model.register_forward_pre_hook(lambda module, input: print(f"forward pre hook (input: {len(input)})"))
model.register_forward_hook(lambda module, input, output: print(f"forward hook (output: {output.shape})"))

y_p = model(x)
# forward pre hook (input: 1)
# forward hook (output: torch.Size([7, 1, 1]))

またreluはプロパティとして定義していない点にも注目です。学習可能なパラメータを持たない層 (relu, tanhやbatch_normなど) は状態の無いただの関数なので、プロパティで持つ必要がありません。

torch以下にはそういった関数形式のインタフェースで層が提供されています。torch.reluもその1つです。モジュールクラスではパラメータをプロパティで持ちますが、関数の場合は引数で重みなども渡します。
nn.ReLUクラスとしても提供されていますので、何度も同じ関数を挟むのであればプロパティとして持って再利用するほうが良いかもしれません。このあたりは好みの問題です。

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.linear_1 = nn.Linear(1, 5)
        self.linear_2 = nn.Linear(5, 1)
        # ...
        
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.linear_1(x)
        x = self.relu(x)
        x = self.linear_2(x)
        x = self.relu(x)
        # ...
        return x

Sequential

上のように「入力xを最初のモジュールの入力として、その出力を2番目のモジュールの入力として、...」という単純な数珠つなぎのforwardであれば、nn.Sequentialが便利です。
Sequentialクラスは上の数珠つなぎのforwardが既に定義されたモジュールなので、各層のモジュールオブジェクトを順番通りに渡すだけでモジュールを作成できます。

model = nn.Sequential(
    nn.Linear(1, 5),
    nn.ReLU(),
    nn.Linear(5, 1)
)

y_p = model(x)

まとめ

今回はPyTorchでネットワークを作りました。

• PyTorchの層はnn.Moduleのサブクラスで実現されており、自分でネットワークを作る場合もnn.Moduleの継承クラスで実装する
• Propertyオブジェクトとしてプロパティに持たせることで、学習可能なパラメータとして認識される
• Moduleサブクラスオブジェクトをプロパティに持つことで、そのプロパティに対しても再帰的にパラメータが探索される

参考

Deep Learning with PyTorch | PyTorch

• 5章を主に参考にしました