PyTorchにおけるカスタムレイヤーの実装

ディープラーニングのモデルを実装する際に用いるライブラリとして、PyTorchを選択する人は多いでしょう。

nn.Linear や nn.Conv2d など、多くのレイヤーが高レベルAPIとして用意されているため、ちょっとしたモデルならばすぐに実装できてしまいますし、複雑なモデルを実装する際も、そのアーキテクチャにのみ集中することができます。

一方、PyTorchで用意されていないようなレイヤーが必要になる場合はどのように実装すれば良いでしょうか？特に難しいことはないのですが、意外と知られていないことが多いので、簡単な例で見ていきましょう。

今回は、全結合層（＋活性化層）を Dense クラスとして実装してみます。その際、レイヤーのパラメータは重みとバイアスになりますが、これをどのように定義するかがポイントになります。以下がコードです。

import numpy as np
import torch
import torch.nn as nn


class Dense(nn.Module):
    def __init__(self, input_dim, output_dim,
                 activation=lambda x: x):
        '''
        引数:
            input_dim: 入力次元
            output_dim: 出力次元
            activation: 活性化関数
        パラメータ:
            W: 重み
            b: バイアス
        '''
        super().__init__()
        self.W = \
            nn.Parameter(torch.Tensor(
                np.random.normal(size=(input_dim,
                                       output_dim))))
        self.b = \
            nn.Parameter(torch.Tensor(
                np.zeros(output_dim)))
        self.activation = activation

    def forward(self, x):
        return self.activation(
            torch.matmul(x, self.W) + self.b)

カスタムレイヤーを定義する場合のポイントは次の３点でしょう。

nn.Module を継承したクラスを定義する
レイヤーのパラメータは nn.Parameter を用いて定義する
順伝播は forward メソッドを定義する

とは言うものの、１つ目と３つ目はモデルを定義するときと変わりませんので、実質カスタムレイヤーを定義するときのみ意識すべきポイントは２つ目になります。nn.Parameter を用いることで、レイヤーのパラメータとして定義することができます。

ここでは重み W の初期値として標準正規分布に従う乱数を発生させていますが、nn.Parameterの引数としては通常のPyTorchのテンソル型となるので、初期化処理などは自由に行うことができます（Xavierの初期化やHeの初期化を行うのが望ましいでしょう）。

今回定義した Dense レイヤーは、他の nn.Linear などと同様、モデルの定義に用いることができます。例えば２クラス分類のロジスティック回帰のモデルを LogisticRegression クラスとして定義してみると、次のように書くことができます。

class LogisticRegression(nn.Module):
    '''
    ロジスティック回帰
    '''
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.layer = Dense(input_dim, output_dim,
                           activation=torch.sigmoid)

    def forward(self, x):
        return self.layer(x)

試しに、ORのデータを用いて学習してみましょう。

np.random.seed(1234)
torch.manual_seed(1234)
device = torch.device('cuda'
                      if torch.cuda.is_available()
                      else 'cpu')

'''
データの読み込み
'''
x = torch.Tensor([[0., 0.],
                  [0., 1.],
                  [1., 0.],
                  [1., 1.]]).to(device)
t = torch.Tensor([[0], [1], [1], [1]]).to(device)
'''
モデルの構築
'''
model = LogisticRegression(2, 1).to(device)

'''
モデルの学習
'''
criterion = nn.BCELoss()
optimizer = optimizers.SGD(model.parameters(), lr=0.1)


def compute_loss(label, pred):
    return criterion(pred, label)


def train_step(x, t):
    model.train()
    preds = model(x)
    loss = compute_loss(t, preds)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss, preds


epochs = 1000
for epoch in range(epochs):
    train_loss = 0.
    train_acc = 0.

    loss, preds = train_step(x, t)
    train_loss = loss.item()

    if epoch % 100 == 99 or epoch == epochs - 1:
        print('epoch: {}, loss: {:.3}'.format(
            epoch + 1,
            train_loss,
        ))

上記を実行してみると、下記の結果が得られ、学習が進んでいることが確認できます。

epoch: 100, loss: 0.382
epoch: 200, loss: 0.29
epoch: 300, loss: 0.232
epoch: 400, loss: 0.192
epoch: 500, loss: 0.164
epoch: 600, loss: 0.142
epoch: 700, loss: 0.126
epoch: 800, loss: 0.112
epoch: 900, loss: 0.101
epoch: 1000, loss: 0.0923

それぞれのデータの予測確率も出力してみましょう。

'''
モデルの評価
'''
model.eval()
for (x_, t_) in zip(x, t):
    print('{} => {:.3f}'.format(x_.cpu().detach().numpy(),
                                model(x_).cpu().detach().numpy()[0]))
[0. 0.] => 0.192
[0. 1.] => 0.922
[1. 0.] => 0.929
[1. 1.] => 0.998

これで、カスタムレイヤーを用いてモデルを定義することができました！今回は非常に簡単なレイヤーを定義しましたが、応用的なレイヤーも、実装すべきことは変わりません。何がレイヤーのパラメータなのかをきちんと踏まえ、nn.Parameter を用いて初期化するだけです。ぜひ自分でも実装してみましょう！