실습: 선형 계층

직접 구현하기

선형 계층은 행렬 곱 연산과 브로드캐스팅 덧셈 연산으로 이루어져 있습니다. 선형 계층의 파라미터 행렬 $W$ 가 행렬 곱 연산에 활용 될 것이고, 파라미터 벡터 $b$ 가 브로드캐스팅 덧셈 연산에 활용 될 것입니다. 앞서 배운 파이토치 연산 방법이므로 그대로 다시 구현해보도록 하겠습니다.

W = torch.FloatTensor([[1, 2],
                       [3, 4],
                       [5, 6]])
b = torch.FloatTensor([2, 2])

보이는 바와 같이 $3\times2$ 크기의 행렬 $W$ 와 $2$ 개의 요소를 갖는 벡터 $b$ 를 선언하였습니다. 그럼 이 텐서들을 파라미터로 삼아, 선형 계층 함수를 구성해볼 수 있습니다.

def linear(x, W, b):
    y = torch.matmul(x, W) + b
    
    return y

아주 쉽죠. matmul 함수와 브로드캐스팅을 이용한 + 연산자만 있으면 됩니다. 그럼 다음과 같이 임의의 텐서를 만들어 함수를 통과시키면, 선형 계층을 통과시키는 것과 같습니다.

x = torch.FloatTensor(4, 3)

3개의 요소를 갖는 4개의 샘플을 행렬로 나타내면 x와 같이 $4\times3$ 크기의 행렬이 될 것입니다. 이것을 다음 코드와 같이 함수를 활용하여 선형 계층을 통과시킬 수 있습니다.

y = linear(x, W, b)
print(y.size())

하지만 이 방법은 파이토치 입장에서 제대로 된 계층^layer로 취급되지 않습니다. 그럼 제대로 계층을 만드는 방법도 살펴보겠습니다.

torch.nn.Module 클래스 상속 받기

파이토치에는 nn^{neural networks} 패키지가 있고, 내부에는 미리 정의된 많은 신경망들이 들어있습니다. 그리고 그 신경망들은 torch.nn.Module 이라는 추상클래스를 상속받아 정의되어 있습니다. 우리도 이 추상클래스를 상속받아 선형 계층을 구현할 수 있습니다. 그에 앞서 torch.nn 패키지를 불러옵니다.

import torch.nn as nn

그리고 nn.Module을 상속받은 MyLinear라는 클래스를 정의합니다. nn.Module을 상속받은 클래스는 보통 2개의 메서드^method __init__ 과 forward 를 오버라이드^override합니다. __init__ 함수는 계층 내부에서 필요한 변수를 미리 선언하는 부분이며, 심지어는 또 다른 계층(nn.Module을 상속받은 클래스의 객체)을 소유하도록 할 수도 있습니다. forward 함수는 계층을 통과하는데 필요한 계산을 수행하도록 구현하는 부분입니다.

class MyLinear(nn.Module):

    def __init__(self, input_dim=3, output_dim=2):
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        super().__init__()
        
        self.W = torch.FloatTensor(input_dim, output_dim)
        self.b = torch.FloatTensor(output_dim)

    # You should override 'forward' method to implement detail.
    # The input arguments and outputs can be designed as you wish.
    def forward(self, x):
        # |x| = (batch_size, input_dim)
        y = torch.matmul(x, self.W) + self.b
        # |y| = (batch_size, input_dim) * (input_dim, output_dim)
        #     = (batch_size, output_dim)
        
        return y

forward 함수 내부에 주석으로 계산되고 있는 텐서들의 모양을 적어놓은 것을 볼 수 있습니다. 이처럼 독자분들도 구현할 때 주석을 통해 텐서의 모양을 미리 적어놓으면, 추후에 코드를 파악하거나 디버깅 할 때 훨씬 편리할 수 있습니다. 앞서와 같이 정의된 클래스를 이제 생성하여 사용할 수 있습니다.

linear = MyLinear(3, 2)

y = linear(x)

여기서 중요한 점은 forward 함수를 따로 호출하지 않고, 객체명에 바로 괄호를 열어 텐서 x를 인수로 넘겨주었다는 것입니다. 이처럼 nn.Module의 상속받은 객체는 __call__ 함수와 forward가 맵핑되어 있어, forward를 직접 부를 필요가 없습니다. 사실은 forward 호출 앞뒤로 파이토치 내부에서 추가적으로 호출하는 함수들이 따로 있기 때문에, 사용자가 직접 forward 함수를 호출하는 것은 권장되지 않습니다.

여기까지 우리는 nn.Module을 상속받아 선형 계층을 구성하여 보았습니다. 하지만 이 방법도 아직은 제대로 된 방법이 아닙니다. 왜냐하면 파이토치 입장에서는 비록 MyLinear라는 클래스의 계층으로 인식하고 계산도 수행하지만, 내부에 학습할 수 있는 파라미터는 없는것으로 인식하기 때문입니다.^[1] 예를 들어 다음의 코드를 실행하면 아무것도 출력되지 않을 것입니다.

for p in linear.parameters():
    print(p)

[1]: 비록 우리는 아직 학습하는 방법에 대해 배우지는 않았습니다만..

올바른 방법: nn.Parameter 활용하기

제대로된 방법은 W와 b를 파이토치에서 학습 가능하도록 인식할 수 있는 파라미터로 만들어주어야 합니다. 이것은 torch.nn.Parameter 클래스^[2]를 활용하면 쉽게 가능합니다. 다음의 코드와 같이 파이토치 텐서 선언 이후에 nn.Parameter로 감싸주면 됩니다.

class MyLinear(nn.Module):

    def __init__(self, input_dim=3, output_dim=2):
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        super().__init__()
        
        self.W = nn.Parameter(torch.FloatTensor(input_dim, output_dim))
        self.b = nn.Parameter(torch.FloatTensor(output_dim))
        
    def forward(self, x):
        # |x| = (batch_size, input_dim)
        y = torch.matmul(x, self.W) + self.b
        # |y| = (batch_size, input_dim) * (input_dim, output_dim)
        #     = (batch_size, output_dim)
        
        return y

그럼 다음과 같이 파라미터를 출력하도록 했을 때, 파라미터가 정상적으로 출력되는 것을 볼 수 있습니다.

>>> for p in linear.parameters():
...     print(p)
Parameter containing:
tensor([[ 0.0000e+00, -0.0000e+00],
        [ 9.1002e+31, -2.8586e-42],
        [ 8.4078e-45,  0.0000e+00]], requires_grad=True)
Parameter containing:
tensor([4.7428e+30, 7.1429e+31], requires_grad=True)

[2]: https://pytorch.org/docs/stable/nn.html#torch.nn.Parameter

nn.Linear 활용하기

사실은 위의 복잡한 방법들은 모두 필요 없고, torch.nn에 미리 정의된 선형 계층을 불러다 쓰면 매우 간단합니다. 다음의 코드는 nn.Linear 를 통해 선형 계층을 활용하는 모습입니다.

linear = nn.Linear(3, 2)

y = linear(x)

그리고 마찬가지로 파라미터도 잘 갖고 있는 것을 볼 수 있습니다.

>>> for p in linear.parameters():
...     print(p)
Parameter containing:
tensor([[-0.3386, -0.5355, -0.4991],
        [ 0.1993,  0.4776,  0.1894]], requires_grad=True)
Parameter containing:
tensor([0.1819, 0.0941], requires_grad=True)

또한 앞서 말한 대로 nn.Module을 상속받아 정의한 나만의 계층 클래스는 내부에 nn.Module의 하위 클래스를 소유할 수 있습니다.

class MyLinear(nn.Module):

    def __init__(self, input_dim=3, output_dim=2):
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        super().__init__()
        
        self.linear = nn.Linear(input_dim, output_dim)
        
    def forward(self, x):
        # |x| = (batch_size, input_dim)
        y = self.linear(x)
        # |y| = (batch_size, output_dim)
        
        return y

앞서의 코드는 nn.Module을 상속받아 MyLinear 클래스를 정의하고 있는데, __init__ 함수 내부에는 nn.Linear를 선언하여 self.linear에 저장하는 모습입니다. 그리고 forward 함수에서는 self.linear에 텐서 x를 통과시킵니다. 즉, 이 코드도 선형 계층을 구현한 것이라 볼 수 있습니다.^[3]

[3]: 다만 바깥의 MyLinear는 빈 껍데기에 불과합니다.