PyTorch torch.nn.Module 関数

PyTorch torch.nn 参考手册PyTorch torch.nn リファレンスマニュアル

torch.nn.ModulePyTorchにおけるすべてのニューラルネットワークモジュールの基底クラスです。

すべてのカスタムネットワークモデルはこのクラスを継承する必要があります。このクラスはパラメータ管理、デバイス移行、モデル保存などの機能を提供します。

クラス定義

torch.nn.Module

主要属性

  • parameters(): モデルのすべての学習可能なパラメータを返します
  • named_parameters(): パラメータ名と値のイテレータを返します
  • children(): モデルのサブモジュールを返します
  • named_children(): サブモジュール名とモジュールのイテレータを返します
  • modules(): すべてのモジュールを返します
  • state_dict(): すべてのパラメータを含む辞書を返します

使用例

例1: カスタムモジュールの作成

nn.Moduleを継承してカスタムネットワークを作成：

インスタンス

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNet, self).__init__()
        # ネットワーク層を定義
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # フォワード伝播を定義
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# モデルをインスタンス化
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)

# テスト
x = torch.randn(32, 784)
output = model(x)

print("モデル構造:")
print(model)
print("n入力形状:", x.shape)
print("出力形状:", output.shape)

例2: パラメータの管理

モデルのパラメータにアクセスして管理：

インスタンス

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.bn1 = nn.BatchNorm2d(16)
        self.fc = nn.Linear(16 * 14 * 14, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = Net()

# パラメータ数を集計
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print("総パラメータ数:", total_params)
print("トレーニング可能なパラメータ数:", trainable_params)

# 特定のパラメータにアクセス
print("nconv1 重み形状:", model.conv1.weight.shape)
print("fc バイアス形状:", model.fc.bias.shape)

例3: モデルの保存と読み込み

モデルの保存と読み込み：

インスタンス

import torch
import torch.nn as nn
import tempfile
import os

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

model = Net()

# モデル全体を保存
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
    torch.save(model, f.name)
    path_full = f.name

# state_dictを保存（推奨方法）
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
    torch.save(model.state_dict(), f.name)
    path_state = f.name

# モデルを読み込み
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()

# 読み込んだモデルをテスト
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)

print("元の出力:", output1[0].tolist())
print("読み込み後の出力:", output2[0].tolist())

# クリーンアップ
os.remove(path_full)
os.remove(path_state)

例4: デバイス移行

異なるデバイス間でモデルを移行：

インスタンス

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

model = Net()

# 現在のデバイスを確認
print("パラメータのデバイス:", model.fc.weight.device)

# GPUに移動（可能な場合）
if torch.cuda.is_available():
    model = model.cuda()
    print("GPUに移動後:", model.fc.weight.device)

# CPUに戻す
model = model.cpu()
print("CPUに戻した後:", model.fc.weight.device)

例5: applyを使用した初期化

applyを使用して再帰的に初期化：

インスタンス

import torch
import torch.nn as nn

def init_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.xavier_uniform_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 5)

    def forward(self, x):
        return self.fc2(self.fc1(x))

model = Net()
model.apply(init_weights)

print("fc1 重み:", model.fc1.weight[0, :3].tolist())
print("fc2 重み:", model.fc2.weight[0, :3].tolist())

例6: 複雑なネットワーク構造

分岐のあるネットワークを構築：

インスタンス

import torch
import torch.nn as nn

class BranchNet(nn.Module):
    def __init__(self):
        super(BranchNet, self).__init__()
        # バックボーン
        self.shared = nn.Linear(10, 20)
        # 分岐
        self.branch_a = nn.Linear(20, 5)
        self.branch_b = nn.Linear(20, 3)

    def forward(self, x):
        feat = self.shared(x)
        out_a = self.branch_a(feat)
        out_b = self.branch_b(feat)
        return out_a, out_b

model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)

print("分岐Aの出力:", out_a.shape)
print("分岐Bの出力:", out_b.shape)

よくある質問

Q1: なぜ super().__init__() を呼び出す必要があるのか？

これは親クラスの初期化メソッドを呼び出し、モジュールがパラメータを正しく登録することを保証します。

<h3 Q2: モデル構造を確認するには？

直接 print(model) するか、torchvision の summary 関数を使用します。

<h3 Q3: parameters() と modules() の違いは？

parameters() は学習可能なパラメータを返し、modules() はすべてのモジュールを返します。

使用シーン

nn.Moduleすべてのカスタムニューラルネットワークを構築するための基礎であり、主な使用シーンは次のとおりです：

  • ネットワーク構造の定義: あらゆるカスタムニューラルネットワーク
  • モデル管理: 保存、読み込み、移行
  • パラメータ管理: オプティマイザーの更新、勾配計算

ヒント：すべてのカスタムModuleは forward メソッドを実装する必要があります。

PyTorch torch.nn 参考手册PyTorch torch.nn リファレンスマニュアル

その他の拡張