nn.Sequentialの挙動を実装して理解する。

この記事では、nn.Sequentialの基本的な概念から始め、自作のCustomSequentialクラスを通してその挙動を解説します。

1. nn.Sequentialとは？

nn.Sequentialは、PyTorchでネットワークを構築するためのシンプルで便利なツールです。これは順番にモジュールを適用することができ、モデルの構築を簡略化します。まず、基本的な使い方を見てみましょう。

import torch.nn as nn

# nn.Sequentialを使った例
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

上記の例では、3つの層が順番に適用されるモデルが構築されています。

2. CustomSequentialクラスの作成

まずは、nn.Sequentialと同様の動作を持つ簡単なカスタムクラスCustomSequentialを作成しましょう。このクラスでは、与えられたモジュールを順番に適用するforwardメソッドを持っています。

class CustomSequential(nn.Module):
    _modules: dict  # 子モジュールを格納する辞書

    def __init__(self, *args: nn.Module):
        super(CustomSequential, self).__init__()
        
        self._modules = {}

        # argsに渡された層を順に追加
        for idx, layer in enumerate(args):
            self.add_module(f'layer_{idx}', layer)

    def add_module(self, name: str, module: nn.Module) -> None:
        # 子モジュールを追加
        if not isinstance(module, nn.Module):
            raise ValueError(f"{module} is not a Module")
        self._modules[name] = module

    def children(self) -> torch.nn.ModuleDict:
        # 子モジュールを返す
        for name, module in self._modules.items():
            yield module

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 子モジュールを順に適用
        for layer in self.children():
            x = layer(x)
        return x

このクラスはnn.Sequentialと同じように使用でき、カスタムモジュールを順番に適用します。 add_moduleやchildrenなどのメソッドはnn.Module継承時に使えますが、ここでは分かりやすく簡単に実装しています。

3. カスタムクラスの利用

さっそく、先ほど作成したCustomSequentialクラスを使用してモデルを構築し、入力データに適用してみましょう。

# カスタムSequentialを使った例
model = CustomSequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

# 入力データの例
input_data = torch.randn((32, 10))

# フォワードパスの実行
output = model(input_data)

これにより、nn.Sequentialと同じようにカスタムモデルが動作します。CustomSequentialクラス内部では、forwardメソッドで子モジュールを順番に適用しています。

4. 挙動の理解

nn.Sequentialは各モジュールを順番に適用し、その出力を次のモジュールの入力として渡します。これにより、層を簡潔に積み重ねることができ、ネットワークの構築が容易になります。