パイトーチでtorch.nnとtorch.nn.functionalの違い 📂機械学習

パイトーチでtorch.nnとtorch.nn.functionalの違い

説明

PyTorchには、多くのニューラルネットワーク関連の関数がtorch.nnとtorch.nn.functionalに同じ名前で含まれています。 nn の関数はニューラルネットワークを関数として返し、 nn.functional の関数はニューラルネットワークそのものです。

例えば nn.MaxPool2d はカーネルサイズを入力として受け取り、プーリング層を返す。

import torch
import torch.nn as nn

pool = nn.MaxPool2d(kernel_size = 2)
# MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)

A = torch.arange(16.).reshape(1, 4, 4)
# tensor([[[ 0.,  1.,  2.,  3.],
#          [ 4.,  5.,  6.,  7.],
#          [ 8.,  9., 10., 11.],
#          [12., 13., 14., 15.]]])

pool(A)
# tensor([[[ 5.,  7.],
#          [13., 15.]]])

一方で nn.functional.MaxPool2d はそのものが2次元マックスプーリングレイヤーです。そのため、この関数はプーリングを適用するテンソルとプーリングの条件をすべて入力として受け取り、実際に入力テンソルをプーリングした結果を返す。

import torch
import torch.nn.functional as F

A = torch.arange(16.).reshape(1, 4, 4)

F.max_pool2d(A, kernel_size=2)
#tensor([[[ 5.,  7.],
#         [13., 15.]]])

つまり nn.MaxPool2d(kernel_size=(n,m)) が返す関数の forward が max_pool2d( ,kernel_size(n,m)) として定義されているわけです。コードを見ると、実際には次のようになっています。

class MaxPool2d(_MaxPoolNd):

    kernel_size: _size_2_t
    stride: _size_2_t
    padding: _size_2_t
    dilation: _size_2_t

    def forward(self, input: Tensor):
        return F.max_pool2d(input, self.kernel_size, self.stride,
                            self.padding, self.dilation, ceil_mode=self.ceil_mode,
                            return_indices=self.return_indices)

パラメータが含まれるレイヤー、例えば線形層であれば、F.linear(input, weight, bias) のようにパラメータも入力として受け取ります。

環境

OS: Windows11
バージョン: Python 3.11.5, torch==2.0.1+cu118