はじめに
みなさん、Pytorchでニューラルネットワークの構築をしたことはありますか？
実装の過程でこのようなコードが出てきますよね。
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
「このnn.Moduleって結局なんなの？」とか「なんでnn.Moduleが必要なの？」みたいに思ったことありませんか？
本記事ではこのnn.Moduleについてその仕組みと役割を深堀りしようと思います。
この記事には生成AIを使用したコンテンツが含まれています。
内容に不正確な表現があるかと思いますが大目に見てください。
環境
- 使用したライブラリ
- Name : torch
- Version : 2.9.1
- Name : torchvision
- Version : 0.24.1
- 生成系AI
- Name : Gemini 3 Pro
内部変数
Moduleクラスの持つ内部変数は20個程度存在し、全てを取り上げるのは面倒くさい分かりにくいのでユーザの使用によって値が変更される変数(要は使用頻度の高い変数)を取り上げます。
- モデル定義時に書き換わるもの
-
_parameters: 学習対象の重み -
_modules: 自身のモデルの中に定義した別のnn.Module -
_buffers: 学習されないが保存すべき状態 -
_non_persistent_buffers_set:state_dictに書き出さないバッファの名前リスト
-
- モードの切替時に書き換わるもの
-
training: 現在が学習モードか推論モードかを示す
-
メソッド
Moduleクラスの持つメソッドは70個以上存在し、全てを取り上げることは無理なのでユーザが扱うものに限定して取り上げます。
- 基本動作・定義
-
__init__: 初期化 -
forward: 順伝播の計算定義、オーバーライドする必要がある -
__call__:インスタンス(x)の形で呼び出される、forwardを自動で呼び出す
-
- モード切替
-
train,eval: 学習モードと推論モードの切替
-
- デバイス・型変換
-
to: CPU/GPUなどのデバイス移動やfloat/halfなどの型変換 -
cuda: GPUを使用する -
cpu: CPUを使用する -
xpu: XPUを使用する
-
- パラメータ取得
-
parameters: パラメータをイテレータとして返す
-
- 保存・読み込み
-
state_dict: 学習済みパラメータを辞書形式で返す -
load_state_dict: パラメータの辞書をモデルに読み込ませる
-
- その他
-
zero_grad: 勾配のリセットを行う
-
つまり、nn.Moduleとは、_modulesによる階層構造の再現と、_parametersによるパラメータの一元管理という2つの機能により、複雑なニューラルネットワークの効率的な構築を実現する基底クラスだと言えます。
なぜnn.Moduleが必要なのか
ここまでで「nn.Moduleとは何か」について触れてきましたが、もう一つの主題である「なぜnn.Moduleが必要なのか」という点には触れていません。
結論を先に述べるとお察しの通り、簡単に実装することができるからなのですが、これで済ませるのは味気ないです。
そこで、具体的にどれくらい簡単になるのかまで見ていきたいと思います。
検証タスクの構成
- データセット : MNIST(手書き文字 0~9)
- ネットワーク構成 : 3層
- 入力層 : 784次元
- 中間層 : 128次元(活性化関数 : ReLU)
- 出力層 : 10次元
- バッチサイズ : 64
-
nn.Moduleを使わない場合import torch import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader # --- 1. データ準備 (共通) --- transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]) train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True) # --- 2. モデル定義 (nn.Moduleなし) --- class ManualNetwork: def __init__(self): # 重みとバイアスを手動で初期化 (Xavier初期化などを簡略化したもの) # requires_grad=True を忘れると学習されない self.W1 = torch.randn(784, 128) * 0.01 self.b1 = torch.zeros(128) self.W2 = torch.randn(128, 10) * 0.01 self.b2 = torch.zeros(10) self.W1.requires_grad = True self.b1.requires_grad = True self.W2.requires_grad = True self.b2.requires_grad = True # 【面倒ポイント1】Optimizerに渡すために、自分でリストを作らないといけない self.params = [self.W1, self.b1, self.W2, self.b2] def forward(self, x): # 行列演算をすべて自分で書く x = x.view(-1, 784) # Flatten x = x @ self.W1 + self.b1 x = F.relu(x) x = x @ self.W2 + self.b2 return x def to(self, device): # 【面倒ポイント2】Tensorの.to()は非破壊的なので、再代入が必要 # さらに、paramsリストの中身も更新しないとOptimizerが古いCPUの変数を参照してしまう self.W1 = self.W1.to(device).requires_grad_(True) self.b1 = self.b1.to(device).requires_grad_(True) self.W2 = self.W2.to(device).requires_grad_(True) self.b2 = self.b2.to(device).requires_grad_(True) # リストも作り直し self.params = [self.W1, self.b1, self.W2, self.b2] # --- 3. 学習ループ --- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ManualNetwork() model.to(device) # 自作のtoメソッドを呼ぶ # パラメータリストを手動で渡す optimizer = torch.optim.SGD(model.params, lr=0.01) print("Start Training (Manual Implementation)...") for epoch in range(1): for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model.forward(data) loss = F.cross_entropy(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print(f'Batch {batch_idx}: Loss {loss.item():.4f}') -
nn.Moduleを使った場合import torch import torch.nn as nn import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader # --- 1. データ準備 (共通) --- transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]) train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True) # --- 2. モデル定義 (nn.Moduleあり) --- class EfficientNetwork(nn.Module): def __init__(self): super().__init__() # 【楽なポイント1】層を定義するだけでパラメータが自動登録される self.layer1 = nn.Linear(784, 128) self.layer2 = nn.Linear(128, 10) def forward(self, x): x = x.view(-1, 784) x = F.relu(self.layer1(x)) return self.layer2(x) # --- 3. 学習ループ --- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = EfficientNetwork() # 【楽なポイント2】一括でデバイス移動（再代入不要） model.to(device) # 【楽なポイント3】parameters()メソッドで全パラメータを自動取得 optimizer = torch.optim.SGD(model.parameters(), lr=0.01) print("Start Training (nn.Module Implementation)...") for epoch in range(1): for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) # forwardではなくインスタンス呼び出し loss = F.cross_entropy(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print(f'Batch {batch_idx}: Loss {loss.item():.4f}')
なぜ楽になるのか
ManualNetwork(以下MN)とEfficientNetwork(以下EN)の決定的な違いとはパラメータリストの作り方です。MNでは9行で記述した処理をENでは2行で記述しています。この差を生み出しているのは、nn.Linearです。
nn.Moduleを継承したこのクラスは、内部で重みやバイアスをtorch.nn.Parameterとして保持しています。ここで重要になるのが、torch.nn.Parameterが変数に代入された時、nn.Moduleの__setattr__メソッドが自動的に呼び出される点です。
この機能によって、変数名とその値が_parametersに辞書形式で自動登録される仕組みになっています。
この_parametersがあることでモデルの運用は非常にシンプルになります。
MNではパラメータの管理に手動でリストを作成していました。
しかし、ENでは、自動で完結しています。
デバイス間の移動に関しても同様です。
このようにnn.Moduleを使用することで面倒な処理や管理がなくなり、記述漏れなどのヒューマンエラーも未然に防がれていることが分かります。
まとめ
本記事では、Pytorchでの実装において必須となるnn.Moduleについて、その内部変数の役割と実装の裏側を深堀りしました。
nn.Moduleとは変数の定義とパラメータ管理を行う自動管理コンテナです。
その役割は以下の3点に集約されます。
- 自動化 : 代入するだけで、内部辞書へ自動で登録する
- 一貫性 : 登録された全パラメータに対して、デバイス間の移動や保存を一括で行う
- 安全性 : 自動化によってヒューマンエラーを排除する