マスジョイ
PR
python3 機械学習・ディープラーニング
pytorch
更新:2024/10/10

【PyTorch】nn.Moduleとモデルの構築と使い方について

ふゅか
ふゅか
PyTorchでモデルを定義するのって、実はnn.Moduleを継承するだけで簡単にできるんだよ！
はるか
はるか
うん、継承してクラスを作る。そしてレイヤーを定義する。
目次

1. nn.Moduleとは

nn.Moduleは、PyTorchにおけるニューラルネットワークモジュールの基底クラスです。レイヤー（層）やモデル全体を定義する際に使用され、パラメータの管理やモデルのシリアライズなど、多くの機能を提供します。

2. 基本的な使い方

ふゅか
ふゅか
クラス定義をする時、nn.Moduleを継承することでモデルの構造を作るんだよ！ここでレイヤーを定義して、ネットワークの骨組みができるのね。
はるか
はるか
つまり、__init__メソッドで使うレイヤーを宣言する。
ふゅか
ふゅか
そうそう！例えば、線形変換のレイヤーを追加するには、nn.Linearって使うの。784次元から128次元に変換するレイヤーを定義できるの！

2.1. クラスの定義と継承

まず、nn.Moduleを継承したクラスを定義します。

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # レイヤーの定義はここで行います

nn.Moduleを継承すると、コンストラクタ__init__メソッド内でレイヤーなどの構成要素を定義できるようになります。次に、具体的にどのようにレイヤーを定義するかを見ていきます。

2.2. レイヤーの定義

モデル内の各レイヤーは、__init__メソッドの中で定義します。ここでは、入力層、ReLU活性化関数、中間層を順に作成しています。例えば、以下のコードでは、784次元の入力を受け取り、128次元に圧縮する線形層（全結合層）と、128次元の出力をさらに10次元に変換するもう一つの線形層を定義しています。

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(128, 10)

2.3. フォワードメソッドの定義

forwardメソッドでは、データがモデルを通過する際の処理を定義します。具体的には、入力データを順に各レイヤーに渡していき、その結果を返します。以下のコードでは、layer1で次元を縮小し、ReLUで非線形性を導入し、最終的にlayer2で10次元に変換する処理を定義しています。

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(128, 10)
    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        return x

2.4. モデルのインスタンス化と使用

モデルを定義したら、次にモデルのインスタンスを作成してデータを流すことができます。以下のコードでは、MyModelをインスタンス化し、ランダムなデータを入力して出力を確認しています。

model = MyModel()
input_data = torch.randn(64, 784)  # バッチサイズ64、入力次元784
output = model(input_data)
print(output.shape)  # torch.Size([64, 10])
ふゅか
ふゅか
モデルをインスタンス化すると、実際にそのモデルでデータを処理できるようになるんだよ！
はるか
はるか
インスタンス化して、実際に使ってみる…という流れね。

3. nn.Moduleの重要なメソッドと属性

3.1. parameters()

モデルのすべての学習可能なパラメータを取得します。オプティマイザにパラメータを渡す際に使用します。

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

上記のコードでは、確率的勾配降下法（SGD）を使用して、学習率0.01でパラメータを最適化しています。

3.2. modules()

モデル内のすべてのサブモジュールをイテレートします。

for module in model.modules():
    print(module)

3.3. train()とeval()

モデルを訓練モードまたは評価モードに設定します。

model.train()  # 訓練モード
model.eval()   # 評価モード

3.4. state_dict()

モデルの現在の状態（パラメータ）を辞書として保存・読み込みします。

# モデルの保存
torch.save(model.state_dict(), 'model.pth')

# モデルの読み込み
model.load_state_dict(torch.load('model.pth'))

4. ポイントとまとめ

  • 初期化の注意：super().__init__()は必ず呼び出してください。これにより、親クラスnn.Moduleの初期化処理が行われます。
  • デバイスの指定：モデルとデータを同じデバイス（CPUやGPU）に配置する必要があります。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
input_data = input_data.to(device)
  • カスタムレイヤーの作成：nn.Moduleを継承して独自のレイヤーを作成できます。

ホームに戻る