파이토치에서 torch.nn과 torch.nn.functional의 차이 📂머신러닝

파이토치에서 torch.nn과 torch.nn.functional의 차이

설명

파이토치에는 신경망과 관련된 많은 함수들이 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
Version: Python 3.11.5, torch==2.0.1+cu118