本文へスキップ
BecomeCoder

PyTorchコース · 第4章 PyTorch入門 ― 自動微分と訓練ループ · レッスン13

モデルを書く ― nn.Module

ローカル実施

導入

第3章では W1, b1, W2, b2 を手で用意し、順伝播も自分で書きました。PyTorch では nn.Module継承して、ネットワークの構造を宣言的に書くだけで済みます。

説明

nn.Linear(入力数, 出力数) が「重み付き和+バイアス」の層1つ。重みの初期化も PyTorch がやってくれます。第3章の XOR ネット(入力2 → 隠れ4 → 出力1)は、こう書けます。

import torch
import torch.nn as nn

class XORNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 4)   # 入力2 → 隠れ4
        self.fc2 = nn.Linear(4, 1)   # 隠れ4 → 出力1

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

model = XORNet()
print(model)

forward に順伝播の流れを書くだけ。重み(W1,W2...)は nn.Linear の中に自動で用意され、model.parameters() でまとめて取り出せます。

flowchart LR
  subgraph XORNet
    fc1["nn.Linear(2,4)"] --> s1["sigmoid"] --> fc2["nn.Linear(4,1)"] --> s2["sigmoid"]
  end
  x["入力"] --> fc1
  s2 --> y["出力"]

補足:層を増やすのも一行

深いネットワークにしたければ層を足すだけです。

self.fc1 = nn.Linear(2, 16)
self.fc2 = nn.Linear(16, 16)
self.fc3 = nn.Linear(16, 1)

numpy で W1, W2, W3... を手管理していた煩雑さが消えます。

まとめ

  • nn.Module を継承し、__init__ で層を、forward で順伝播を書く。
  • nn.Linear(in, out) が「重み付き和+バイアス」1層。重みの用意・初期化は自動。
  • 層を増やすのは1行。第3章の手作りネットが、宣言的で読みやすいコードになる。