Python
PyTorch

[Pytorch]自作レイヤーの作り方

初めに

Pytorchを使った自作レイヤー(自作モジュール)の作り方について説明します。

今回説明するのは、様々なレイヤーを組み合わせることによってモデルを構築する方法ではなく、モデルの構築に使用されるレイヤーの作り方に関する説明です。

環境

Python 3
Pytorch 0.4

参照

https://pytorch.org/docs/stable/notes/extending.html#extending-torch-nn

自作レイヤーの作り方

Pytorchをimportします。

import torch
from torch import tensor
import torch.optim as optim
from torch import autograd
import torch.nn as nn
import torch.nn.functional as F

学習パラメータの作り方

tensorをnn.Parameterで囲むことで、学習パラメータにすることができます。

weight = nn.Parameter(tensor(output_features, input_features))

ここで、nn.Parameterで囲むtensorは、学習パラメータの初期値となり、種々の初期値を定義することが可能です。

weight_randn = nn.Parameter(torch.randn(output_features, input_features))
weight_ones = nn.Parameter(torch.ones(output_features, input_features))

このようにnn.Parameterで囲むtensorは自由に決めることができます。
今日では、様々な初期化の手法が報告されていますが、それらを実装する際に参考になるかもしれません。

自作レイヤーを作る

  • nn.Moduleを継承している。
  • __ init__ を実装している。(必須ではない)
  • forward()を実装している。

これら3つの要件を満たすクラスを実装することが必要です。

class MyLayer(nn.Module):
    def __init__(self, input_features, output_features):
        super(MyLayer, self).__init__()
        self.input_features = input_features
        self.output_features = output_features

        self.weight = nn.Parameter(torch.randn(input_features, output_features))

    def forward(self, input):
        return torch.matmul(input, self.weight)

このように、
__ init__ には、学習パラメータの入力や出力の数(特徴の数)、初期化の方法について定義します。

forward() には、この自作レイヤーが行う操作(計算)について定義します。上記は、入力された特徴に対して、学習パラメータを掛け合わせたものを返しています。

以上、自作レイヤーの作り方です。

自作レイヤーを使用する。

layer = MyLayer(input_features=7, output_features=2)
inputs = torch.randn(6,7)
result = layer(inputs)

#resultの和を用いて微分する。
result.sum().backward()
layer.weight.grad

実行すると、インスタンス化されたlayerの持つweightに、勾配が保持されているのが確認できます。
したがって、layerは学習パラメータをもつモジュールとして機能します。

通常実装されているレイヤーと同様に、モデルの構築に利用することができます。

モデルの構築方法については、分かりやすい説明がQiitaに存在しているので確認してみてください。

以上で終わりです。
ご指摘等ございましたら、ぜひお願いします。