パイトーチで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