[PyTorch/torchvision] ネットワークを設定する (3)

ゼロからネットワークを作成

既存のネットワークではなく、自分でゼロからネットワークを設定してみます。ネットワークのコードは以下を使って、実装方法を確認してみました。
ちなみに、このコードはMNISTという0~9の数字 (28×28ピクセルの画像) を分類するためのモデルとして設計されています。

Build the Neural Network – PyTorch Tutorials 1.12.0+cu102 documentation

ネットワークを構成するためのモジュール類はtorch.nnに入っています。
まずは、必要モジュールを一式インポートします。

import torch
from torch import nn

ネットワークの定義

ネットワークの定義に必要な各要素はnn.Moduleのサブクラスに含まれているため、nn.Moduleを継承したクラスとして設定します。

設定するモデルは以下の2つの関数が含まれている必要があります。

• __init__ : 初期化
• forward : feed forward処理を実施

クラスの外部とやり取りをするのはこの2つだけですので、内部的な処理を行う関数は別途、任意で作成しても問題ありません。

上記リンクで定義されているネットワークは以下のシンプルなものです。全結合層とReLUから構成されていて、最終出力のlogitsを後処理でSoftmaxにかけることで、入力画像が0~9のそれぞれの数字である確率が求められることになります。

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()

        self.flatten = nn.Flatten()  # 2次元画像を1次元に変換

        # 3つの全結合層 (ReLUを間に含む) からなるネットワーク
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512), 
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()

定義後のモデルの中身はこんな感じです。

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

nn.Sequential()モジュール

ここで、nn.Sequential()というものが出てきました。これは、処理順序を定義したコンテナです。nn.Sequential()にレイヤを入れていくことで、格納された順番に処理が進んでいくことになります。

上の例だと、nn.Linear(28*28, 512)のアウトプットが、nn.ReLU()のインプットとなる、というイメージです。
各レイヤはスライスで取得することができます。

print(model.linear_relu_stack[0]) 

# Linear(in_features=784, out_features=512, bias=True)

OrderedDict

順序付きの辞書であるOrderedDictを使うことで、スライスだけでなく、名前指定でもレイヤを取得できるようになります。

from collections import OrderedDict  # OrderedDictをインポート

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(OrderedDict([
            ('fc1', nn.Linear(28*28, 512)),
            ('relu1', nn.ReLU()),
            ('fc2', nn.Linear(512, 512)),
            ('relu2', nn.ReLU()),
            ('fc3', nn.Linear(512, 10)),
        ]))

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()

こんな感じでレイヤに名前が付きました。

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (fc1): Linear(in_features=784, out_features=512, bias=True)
    (relu1): ReLU()
    (fc2): Linear(in_features=512, out_features=512, bias=True)
    (relu2): ReLU()
    (fc3): Linear(in_features=512, out_features=10, bias=True)
  )
)

OrderedDictを使って設定した場合、レイヤは以下の2パターンで呼び出せます。

print(model.linear_relu_stack[0])   # スライス指定
print(model.linear_relu_stack.fc1)  # 名前指定

# どちらの場合も
# Linear(in_features=784, out_features=512, bias=True)