PyTorchでニューラルネットワークを構築：torch.nn.ModuleList入門

PyTorchのニューラルネットワークにおけるtorch.nn.ModuleList

従来のリストとの違い

従来のPythonリストと比較すると、torch.nn.ModuleListには以下の利点があります。

1. モジュールの登録と管理を容易にする: torch.nn.ModuleListは、torch.nn.Moduleを継承したモジュールのみを追加できます。これは、誤ったモジュールを追加してエラーが発生するのを防ぎ、コードの安全性と信頼性を向上させます。
2. 自動登録: torch.nn.ModuleListに追加されたモジュールは、自動的にネットワークのparameters()とmodules()属性に登録されます。これは、モデルのパラメータの更新やモジュールのアクセスを容易にする便利な機能です。
3. シームレスなシリアル化とデシリアライズ: torch.nn.ModuleListは、モデルのシリアル化とデシリアライズにおいてもシームレスに動作します。これは、モデルを保存して再読み込みしたり、異なる環境に移行したりする際に役立ちます。

使用例

torch.nn.ModuleListは、様々なニューラルネットワークアーキテクチャで使用できます。以下に、簡単な例を紹介します。

import torch
from torch import nn

class MyModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = nn.ModuleList([
      nn.Linear(10, 20),
      nn.ReLU(),
      nn.Linear(20, 10)
    ])

  def forward(self, x):
    for layer in self.layers:
      x = layer(x)
    return x

model = MyModel()
print(model)

# 出力:
# MyModel(
#   (layers): ModuleList(
#     (0): Linear(in_features=10, out_features=20, bias=True)
#     (1): ReLU()
#     (2): Linear(in_features=20, out_features=10, bias=True)
#   )
# )

この例では、MyModelというクラスが定義されています。このクラスは、3つの層からなるニューラルネットワークを表しており、torch.nn.ModuleListを使用して層を管理しています。

torch.nn.ModuleListは、PyTorchのニューラルネットワークにおいて、順序付きのモジュールのリストを管理するための便利なクラスです。従来のリストと比較すると、モジュールの登録と管理を容易にする、自動登録、シームレスなシリアル化とデシリアライズといった利点があります。

PyTorchのtorch.nn.ModuleListを使用したサンプルコード

MLP (Multi-Layer Perceptron)

import torch
from torch import nn

class MLP(nn.Module):
  def __init__(self, input_dim, output_dim, hidden_dim=100):
    super().__init__()
    self.layers = nn.ModuleList([
      nn.Linear(input_dim, hidden_dim),
      nn.ReLU(),
      nn.Linear(hidden_dim, output_dim)
    ])

  def forward(self, x):
    for layer in self.layers:
      x = layer(x)
    return x

# サンプルコード
input_dim = 10
output_dim = 3

model = MLP(input_dim, output_dim)
x = torch.randn(10, input_dim)
y = model(x)
print(y.shape) # 出力: (10, 3)

CNN (Convolutional Neural Network)

import torch
from torch import nn

class CNN(nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = nn.ModuleList([
      nn.Conv2d(1, 32, 3, 1, 1),
      nn.ReLU(),
      nn.MaxPool2d(2, 2),
      nn.Conv2d(32, 64, 3, 1, 1),
      nn.ReLU(),
      nn.MaxPool2d(2, 2),
      nn.Flatten(),
      nn.Linear(64 * 4 * 4, 10)
    ])

  def forward(self, x):
    for layer in self.layers:
      x = layer(x)
    return x

# サンプルコード
x = torch.randn(1, 1, 28, 28)
model = CNN()
y = model(x)
print(y.shape) # 出力: (1, 10)

RNN (Recurrent Neural Network)

import torch
from torch import nn

class RNN(nn.Module):
  def __init__(self, input_dim, output_dim, hidden_dim=100):
    super().__init__()
    self.rnn = nn.RNN(input_dim, hidden_dim, batch_first=True)
    self.fc = nn.Linear(hidden_dim, output_dim)

  def forward(self, x):
    x, _ = self.rnn(x)
    x = x[:, -1, :]
    x = self.fc(x)
    return x

# サンプルコード
input_dim = 10
output_dim = 3
seq_len = 20

x = torch.randn(1, seq_len, input_dim)
model = RNN(input_dim, output_dim)
y = model(x)
print(y.shape) # 出力: (1, 3)

torch.nn.ModuleList 以外の方法

ネストされたリスト

class MyModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = [
      nn.Linear(10, 20),
      nn.ReLU(),
      nn.Linear(20, 10)
    ]

  def forward(self, x):
    for layer in self.layers:
      x = layer(x)
    return x

# サンプルコード
model = MyModel()
print(model)

# 出力:
# MyModel(
#   (layers): [
#     (0): Linear(in_features=10, out_features=20, bias=True)
#     (1): ReLU()
#     (2): Linear(in_features=20, out_features=10, bias=True)
#   ]
# )

この方法はシンプルですが、コードの可読性が低下する可能性があります。

OrderedDict

from collections import OrderedDict

class MyModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = OrderedDict([
      ("linear1", nn.Linear(10, 20)),
      ("relu", nn.ReLU()),
      ("linear2", nn.Linear(20, 10))
    ])

  def forward(self, x):
    for name, layer in self.layers.items():
      x = layer(x)
    return x

# サンプルコード
model = MyModel()
print(model)

# 出力:
# MyModel(
#   (layers): OrderedDict([
#     ('linear1', Linear(in_features=10, out_features=20, bias=True)),
#     ('relu', ReLU()),
#     ('linear2', Linear(in_features=20, out_features=10, bias=True))
#   ])
# )

この方法は、キーと値のペアでモジュールを管理できます。

自作のクラス

上記のいずれの方法もニーズに合わない場合は、自作のクラスを作成することができます。

class MyModuleList(nn.Module):
  def __init__(self, modules):
    super().__init__()
    self.modules = modules

  def forward(self, x):
    for module in self.modules:
      x = module(x)
    return x

# サンプルコード
modules = [
  nn.Linear(10, 20),
  nn.ReLU(),
  nn.Linear(20, 10)
]

model = MyModuleList(modules)
print(model)

# 出力:
# MyModuleList(
#   (modules): [
#     (0): Linear(in_features=10, out_features=20, bias=True)
#     (1): ReLU()
#     (2): Linear(in_features=20, out_features=10, bias=True)
#   ]
# )

この方法は、最も柔軟性がありますが、コード量が増えてしまいます。

最適な方法は、プロジェクトの要件と開発者の好みによって異なります。

• シンプルなプロジェクトの場合は、ネストされたリストで十分です。
• モジュールの名前を付けた