はじめに

みなさん、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点に集約されます。

自動化 : 代入するだけで、内部辞書へ自動で登録する
一貫性 : 登録された全パラメータに対して、デバイス間の移動や保存を一括で行う
安全性 : 自動化によってヒューマンエラーを排除する

参考

Pytorchの公式ドキュメント
 nn.Module GitHub
nn.Linear GitHub