PyTorch 기초 - 모델 정의하기

데이터를 준비했으면, 이제 학습할 딥러닝 모델을 정의하면 됩니다. 이 포스트에서는 torch.nn 모듈을 사용해서 모델을 정의하는 방법에 대해 알아보겠습니다.

모델을 구현하기 위해 여기서 사용하는 모듈들을 불러오겠습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

간단한 모델 정의하기

아주 간단한 모델을 표현하는 클래스를 작성해 보겠습니다.

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(2, 16)
        self.linear2 = nn.Linear(16, 1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x)
        return x

작성한 SimpleModel 클래스는 nn.Module를 상속하고 있습니다. 여기서 구현이 필요한 메서드는 __init__와 forward입니다.

• __init__ 메서드에서는 모델 안에서 사용한 모듈을 정의합니다. 생성자 제일 앞에 super() 함수를 호출하고(그렇지 않으면 에러가 발생합니다), 사용할 모듈들을 인스턴스 변수로 초기화합니다. 인스턴스 변수에는 PyTorch에서 제공하는 컨볼루션, Fully-connected 레이어들뿐만 아니라 nn.Module을 상속해 직접 작성한 모델도 넣을 수 있습니다.

• forward 메서드는 입력 데이터가 주어졌을 때 선언한 인스턴스 변수(모듈)들을 통과해 출력 데이터를 만드는 메서드입니다.

이 모델에서는 Linear 클래스를 사용합니다. 이는 Fully-connected 레이어라고 볼 수 있는데, 정의는 다음과 같습니다.

torch.nn.Linear(in_features, out_features, bias=True, ...)

이 레이어에 필요한 입력 텐서의 크기는 [batch_size, in_features]이고 출력 텐서의 크기는 [batch_size, out_features]로 입출력 모두 2D 텐서라고 볼 수 있습니다.

그러면 정의한 두 모델에 대해 입력 텐서를 순방향 통과시켜 보겠습니다.

x = torch.randn(8, 2)
fc = nn.Linear(2, 32)
simple_model = SimpleModel()

print("fc(x) shape:\n", fc(x).shape)
print("simple_model(x) shape:\n", simple_model(x).shape) 

fc(x) shape:
 torch.Size([8, 32])
simple_model(x) shape:
 torch.Size([8, 1])

위 코드와 같이 정의한 모델을 함수처럼 입력 텐서를 매개변수로 넣으면, 연산을 수행한 결과 텐서가 나오게 됩니다. 이는 forward 메서드를 수행한 것으로 보시면 됩니다.

nn.Sequential과 nn.ModuleList

위 SimpleModel에서 fully-connected 레이어의 수가 많아질수록, 통과한 데이터를 저장하고, 그 다음 레이어에 넣는 코드를 직접 작성을 해야 한다면 모델을 구현하는데 필요한 줄의 길이가 늘어날 것입니다. 이 과정을 수행할 필요 없이 여러 레이어들을 저장하고 간단하게 모델을 통과하기 위해 nn.Sequential과 nn.ModuleList를 사용할 수 있습니다. 예시를 들어 SimpleModel에 대해 두 클래스들을 각각 사용하면서 같은 역할을 수행하도록 만들어보겠습니다.

class SimpleSequentialModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(2, 16),
            nn.Linear(16, 1)
        )

    def forward(self, x):
        x = self.layers(x)
        return x

nn.Sequential은 생성자에서 정의한 레이어들을 순서대로 통과시킵니다. 여기서 참고할 점은 입력이 리스트가 아니라 여러개의 인수를 두는 점입니다. 정의한 컨테이너를 사용하기 위해서 함수처럼 입력 텐서를 넣으면 바로 레이어들을 통과한 결과를 반환해주기 때문에 편리합니다.

def SimpleSequentialModel2():
    return nn.Sequential(
        nn.Linear(2, 16),
        nn.Linear(16, 1)
    )

참고로 한 개의 nn.Sequential의 결과를 바로 리턴하는 모델을 활용할 때에는, 굳이 모듈 상속 필요 없이 함수로 만들어서 반환하기만 해도 같은 동작을 수행하면서도 코드를 더 짧게 만들 수 있습니다.

class SimpleModuleListModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(2, 16),
            nn.Linear(16, 1)
        ])

    def forward(self, x):
        for linear in self.layers:
            x = linear(x)
        return x

nn.ModuleList은 말 글대로 모듈들을 저장하고 있는 리스트라고 할 수 있습니다. nn.Sequetial과 다른 점은 정의할 때 리스트 형태로 들어간다는 점과, forward 함수에서 각 인덱스에 접근해 해당되는 레이어를 통과하는 점을 명시해야 한다는 점입니다. 그러면 그냥 리스트를 만드는 것과 무엇이 다르냐고 할 수 있는데, nn.Module에서는 리스트에 모듈을 넣을 수 없고, 이를 지키지 않을 경우 모델을 통과할 때 에러를 발생시키므로 꼭 넣어야 합니다. 참고로 리스트와 같이 append, extend, insert 메서드를 지원합니다.

CNN에 자주 사용되는 블록들

이제 본격적으로 CNN을 구성해보고자 합니다. CNN을 구성하기 위해서는 컨볼루션, 배치 정규화, 활성 함수, 풀링 블록 등 많은 연산을 이용하게 됩니다. PyTorch에서는 이를 블록 단위로 제공하고 있습니다. 여기서 이미지가 입력일 때에는 [batch_size, channels, height, width]의 4D 텐서를 사용하는데, 일부 2d가 들어간 블록들은 4D 텐서를 기반으로 연산하도록 설계되었습니다. 이를 염두에 두고 블록들을 종류별로 알아보겠습니다.

먼저 컨볼루션 블록들입니다.

torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, ...)
torch.nn.ConvTranspose2d(...)

이 연산들을 수행하기 위해 먼저 필요한 것은 입력과 출력 채널수입니다. 이는 4D 텐서에서 입력받는 채널의 크기와, 이를 내보낼 출력 채널의 크기를 입력합니다. 예를 들어 RGB 영상의 경우 제일 먼저 통과하는 컨볼루션 레이어의 in_channels는 3이 되겠죠. 그 이후 out_channels로는 직접 지정해줄 수 있습니다. 그 다음에 통과할 연산이 컨볼루션 레이어라면 in_channels에는 이전 레이어의 out_channels를 입력해 주면 됩니다.

그 다음 컨볼루션 커널의 크기, 스트라이드, 패딩 크기를 입력합니다. 컨볼루션 연산에서 스트라이드는 이미지의 크기를 줄이는 역할을 하고, 전치컨볼루션 연산은 이미지 크기를 늘리는 역할을 수행합니다. 패딩은 가장자리에 특정 값을 채워넣는 역할을 수행하는데 이는 컨볼루션을 수행할 때 이미지의 크기가 줄어드는 것을 고려해야 하기 때문입니다. 입출력 이미지의 가로, 세로 크기가 같기 위해 보통 kernel_size // 2의 값을 사용합니다. 참고로 2D 이미지이기 때문에 숫자 대신 2개의 값을 가진 튜플로 가로, 세로의 크기를 다르게 할 수 있습니다.

다음은 풀링 블록들에 대해 살펴보겠습니다.

torch.nn.MaxPool2d(kernel_size, ...)
torch.nn.AvgPool2d(kernel_size, ...)
torch.nn.MaxUnpool2d(kernel_size, ...)
torch.nn.AvgUnpool2d(kernel_size, ...)
torch.nn.AdaptiveMaxPool2d(output_size, return_indices=False)
torch.nn.AdaptiveAvgPool2d(output_size)

일반적인 풀링은 커널 크기를 지정해주면 커널 크기만큼 이미지의 크기가 줄어듭니다. Adaptive 풀링의 경우 이미지 크기를 조절하는 것처럼 결과로 나올 가로, 세로 크기를 튜플 형태로 넣어주면 해당 크게 맞게 풀링을 수행해줍니다.

다음은 딥러닝 모델에서 비선형성에 빠질 수 없는 활성 함수 블록입니다.

torch.nn.ReLU(inplace=False)
torch.nn.LeakyReLU(negative_slope=0.01, inplace=False)
torch.nn.Sigmoid()
torch.nn.Tanh()
torch.nn.Softmax(dim=None)

일반적으로 ReLU, Sigmoid같은 일반적인 활성 함수 모두 구현되어있기 때문에 Sequential 모델에 바로 넣을 수 있습니다. 여기서 inplace라는 옵션을 사용하면 새로운 메모리 할당 없이 입력 텐서에 바로 활성 함수가 적용되어 메모리를 절약할 수 있습니다. 즉 입력 텐서에 덮어씌우기하기 때문에 고려해서 사용하면 됩니다.

마지막으로 배치 정규화 블록들입니다.

torch.nn.BatchNorm2d(num_features, ...)
torch.nn.InstanceNorm2d(num_features, ...)

배치 정규화를 사용하기 위해서는 채널의 수가 필요합니다. 이는 컨볼루션 레이어의 결과 채널 수, 즉 out_channels를 입력해 주면 됩니다.

CNN 모델 정의하기

이제 알아본 블록들을 활용해 CNN 모델을 작성해보겠습니다.

def SimpleCNNLayer(in_features, out_features):
    return nn.Sequential(
        nn.Conv2d(in_features, out_features, 3, 1, 1),
        nn.BatchNorm2d(out_features),
        nn.ReLU(),
        nn.MaxPool2d(2)
    )

def SimpleCNNClassifier():
    return nn.Sequential(
        SimpleCNNLayer(1, 32),
        SimpleCNNLayer(32, 32),
        nn.Flatten(),
        nn.Linear(1568, 64),
        nn.Linear(64, 10)
    )

간단한 CNN 모델을 활용하고자 nn.Sequential을 주로 활용했습니다. SimpleCNNClassifier로 전체적인 모델의 역할을 수행하고, 각 레이어들을 SimpleCNNLayer로 두었습니다. 여기서 SimpleCNNLayer 인수를 초기화할 때 in_features과 out_features를 넘겨주는 방식을 사용하는데, 레이어를 제작하는 방법으로 자주 활용됩니다.

두 SimpleCNNLayer를 지나게 되면 nn.Flatten 블록을 통과하게 됩니다.

torch.nn.Flatten(start_dim=1, end_dim=- 1)

이는 각 배치에 대해 펴주어서 2D 텐서로 만들어주는 역할을 수행합니다. 이 때 위치 정보가 사라진다고도 표현할 수 있습니다. nn.Linear로 최종적으로 판단하게 되는데 여기서 입력으로 들어가는 1568의 경우 두 번째 SimpleCNNLayer를 통과한 크기인 32 x 7 x 7로 볼 수 있습니다.

그러면 임의의 텐서를 넣어서 결과가 나오는지 확인하겠습니다.

x = torch.randn(8, 1, 28, 28)
cnn_model = SimpleCNNClassifier()
print("cnn(x) shape:\n", cnn_model(x).shape)

cnn(x) shape:
 torch.Size([8, 10])

입력에 대해 출력 텐서가 2차원으로 각 인덱스마다 점수를 출력하는 형식으로 잘 나오는 것을 알 수 있습니다.

nn.functional 사용하기

위에서 사용한 nn.Conv2d, nn.MaxPool2d, nn.ReLU같은 블록들은 nn.Module을 상속한 클래스 내의 생성자에서 블록의 객체 정의를 통해 인스턴스 변수로 만들고, forward 함수에서 입력 텐서를 함수처럼 통과시키는 방법을 사용합니다. 이런 블록을 별도로 정의하고 사용하지 않고도 nn.functional 모듈을 활용하면 forward 함수에서 일반적인 함수를 사용하는 것처럼 사용할 수 있습니다.

PyTorch 가이드로 볼 수 있듯이, 컨볼루션을 포함하는 일부 연산을 제공하는 nn.functional 모듈로 일부 연산을 대체할 수 있습니다. 그렇지만 코드의 가독성과 편의성을 고려해서, 일반적으로는 별도의 가중치가 필요하지 않는 활성 함수, 풀링 또는 이미지 크기 조절 등에 사용합니다.

그러면 nn.functional 모듈을 활용한 간단한 CNN 네트워크를 작성해보겠습니다.

class SRCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 9, 1, 4)
        self.conv2 = nn.Conv2d(32, 64, 1, 1, 0)
        self.conv3 = nn.Conv2d(64, 3, 5, 1, 2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.conf3(x)
        return x

모듈을 불러오는 코드에서 nn.functional을 F로 불러왔기 때문에, 이 모듈을 사용할 때 간단하게 F만 명시해 주면 됩니다. 이 예시는 컨볼루션 블록을 통과한 텐서에 대해 ReLU 활성 함수를 적용하기 위해 F.relu 함수를 사용하였습니다. 이런 방식으로 leaky_relu, sigmoid, tanh 함수도 적용할 수 있습니다.

그 다음에 코드에 자주 사용하는 함수는 이미지 크기를 조절하는 풀링과 일반적인 크기 조정 방법(보간법)이 있습니다. 풀링은 max_pool2d, avg_pool2d에 커널 크기 등을 명시해서 사용하면 됩니다. 그리고 이미지 보간같은 경우 블록 형태의 torch.nn.Upsample로도 사용할 수 있지만, 일반적으로 다음 함수를 사용합니다.

torch.nn.functional.interpolate(input, size=None, scale_factor=None, mode='nearest', ...)

여기서 출력될 텐서의 가로 및 세로 크기를 지정하는 size와, 몇 배로 늘릴 것인지를 지정하는 scale_factor 중 하나만 사용하면 됩니다. 이를 사용한 예시를 보이겠습니다.

x = torch.randn(8, 64, 28, 28)
x = F.upsample(x, size=(56, 56), mode='bilinear', align_corners=False)
print("interpolated x shape:", x.shape)

interpolated x shape: torch.Size([8, 64, 56, 56])

가로 및 세로의 길이가 각각 28인 텐서에 대해, bilinear 보간법으로 56으로 늘려주는 역할을 수행합니다. 이는 각 batch, feature map에 대해 수행되며, 결과적으로 가로 및 세로의 크기가 바뀐 것을 확인할 수 있습니다. 참고로 align_corners은 warning을 띄우지 않기 위해 명시한 것입니다.