step40. 브로드캐스트 함수

📢 본 포스팅은 밑바닥부터 시작하는 딥러닝3을 기반으로 작성하였습니다. 배운 내용을 기록하고, 개인적인 공부를 위해 작성하는 포스팅입니다. 자세한 내용은 교재 구매를 강력 추천드립니다.

 

 

이전 단계에서는 DeZero만의 sum 함수를 구현했습니다. 이때 broadcast_to 함수가 사용되었는데, 이번 단계에서 구현해보겠습니다. 순전파에서는 역시 넘파이의 브로드캐스트를 활용하지만, 역전파는 DeZero만의 코드를 구현해내야 합니다. 우선 브로드캐스트 함수와 sum 함수는 상호 의존적인 관계임을 알아야 합니다.

 

 

위의 그림과 같이 코드를 구현하면, broadcast_to의 역전파에는 sum_to가 sum_to의 역전파에는 broadcast_to가 활용되는 것을 확인할 수 있습니다.

class SumTo(Function):
    def __init__(self, shape):
        self.shape = shape
    
    def forward(self, x):
        self.x_shape = x.shape
        y = utils.sum_to(x, self.shape)
        return y

    def backward(self, gy):
        gx = broadcast_to(gy, self.x_shape)
        return gx


def sum_to(x, shape):
    if x.shape == shape:
        return as_variable(x)
    return SumTo(shape)(x)


class BroadcastTo(Function):
    def __init__(self, shape):
        self.shape = shape

    def forward(self, x):
        self.x_shape = x.shape
        y = np.broadcast_to(x, self.shape)
        return y
    
    def backward(self, gy):
        gx = sum_to(gy, self.x_shape)
        return gx
    

def broadcast_to(x, shape):
    if x.shape == shape:
        return as_variable(x)
    return BroadcastTo(shape)(x)

 

다음으로 브로드캐스트 대응이 필요합니다. 넘파이에서는 형상이 다른 다차원 배열끼리의 연산을 가능하게 합니다. DeZero 역시 마찬가지로 형상이 다를 때, 계산 과정에서 브로드캐스트 기능을 대응시키도록 하겠습니다. 다음 코드는 넘파이의 브로드캐스트 기능입니다.

x0 = np.array([1, 2, 3])
x1 = np.array([10])
y = x0 + x1
print(y)  # array([11, 12, 13])

 

핵심은 순전파에서 브로드캐스트가 일어났다면, 역전파에서 브로드캐스트의 역전파가 이루어져야 한다는 점입니다. 앞서 보았던 그림과 같이 broadcast_to 함수의 역전파는 sum_to 함수입니다. 따라서 DeZero의 Add 클래스를 다음과 같이 수정해야 합니다. Mul, Sub, Div 클래스 역시 수정하여야 합니다.

class Add(Function):
    def forward(self, x0, x1):
        self.x0_shape, self.x1_shape = x0.shape, x1.shape
        y = x0 + x1
        return y
    
    def backward(self, gy):
        gx0, gx1 = gy, gy
        if self.x0_shape != self.x1_shape:  # for broadcast
            gx0 = dezero.functions.sum_to(gx0, self.x0_shape)
            gx1 = dezero.functions.sum_to(gx1, self.x1_shape)
        return gx0, gx1

 

마지막으로 다음과 같은 테스트들을 진행할 수 있습니다.

import numpy as np
from dezero import Variable

x0 = Variable(np.array([1, 2, 3]))
x1 = Variable(np.array([10]))
y = x0 + x1
print(y)  # variable([11 12 13])

y.backward()
print(x0.grad)  # variable([1 1 1])
print(x1.grad)  # variable([3])

y = x0 * x1
print(y)  # variable([10 20 30])

x0.cleargrad()
x1.cleargrad()
y.backward()
print(x0.grad)  # variable([10 10 10])
print(x1.grad)  # variable([6])