step18. 메모리 절약 모드

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

 

 

이번 단계에서는 메모리 절약 모드를 구현해보도록 하겠습니다. 우선 필요 없는 미분값들을 삭제하겠습니다. 현재의 DeZero는 미분을 진행하면 모든 변수가 미분 결과를 메모리에 유지합니다. 보통 머신러닝에서는 말단 변수의 미분값만 필요하므로, 중간 변수의 미분값을 제거하는 기능(retain_grad)을 추가합니다.

class Variable:
    ...
    def backward(self, retain_grad=False):
        ...
        while funcs:
            f = funcs.pop()
            gys = [output().grad for output in f.outputs]
            gxs = f.backward(*gys)
            if not isinstance(gxs, tuple):
                gxs = (gxs, )
            
            for x, gx in zip(f.inputs, gxs):
                if x.grad is None:
                    x.grad = gx
                else:
                    x.grad = x.grad + gx

                if x.creator is not None:
                    add_func(x.creator)
            
            if not retain_grad:
                for y in f.outputs:
                    y().grad = None  # y는 약한 참조(weakref)

 

테스트로 다음 코드를 실행해보겠습니다. 중간 변수인 y와 t의 미분값은 삭제되고, 말단 변수인 x0과 x1의 미분값만 유지됩니다. 덕분에 절약된 메모리를 다른 곳에 사용할 수 있게 됩니다.

x0 = Variable(np.array(1.0))
x1 = Variable(np.array(1.0))
t = add(x0, x1)
y = add(x0, t)
y.backward()

print(y.grad, t.grad)  # None None
print(x0.grad, x1.grad)  # 2.0 1.0

 

신경망은 크게 학습과 추론이라는 두 가지 단계로 나뉩니다. 학습 시에는 미분값을 구해야하지만, 추론 시에는 순전파만을 진행하기 때문에 중간 계산 결과를 기록하지 않으면 메모리 사용량을 크게 줄이는 것이 가능합니다. 주로 사용되는 딥러닝 프레임워크인 PyTorch에서는 이 기능을 torch.no_grad()라는 함수 형태로 제공합니다. (매우 자주 쓰입니다.)

 

순전파만 활용할 때를 위해, 역전파 활성 모드와 역전파 비활성 모드를 전환하는 구조를 구축해보겠습니다. 먼저 Config 클래스를 활용하겠습니다. enable_backprop이 True이면 역전파 활성 모드입니다. Config 같은 설정 데이터는 단 한 군데에만 존재하는 것이 좋기 때문에, 인스턴스화하지 않고 클래스 상태로 두겠습니다.

class Config:
    enable_backprop = True

 

이제 Function에서 Config 클래스를 참조할 수 있도록 하겠습니다. 역전파 시에 필요한 세대와 계산들의 연결을 만들어내는 output.set_creator(self)등의 기능을 if문 안에 넣었습니다.

class Function:
    def __call__(self, *inputs):
        ...
        if Config.enable_backprop:
            self.generation = max([x.generation for x in inputs])  # 세대 설정
            for output in outputs:
                output.set_creator(self)  # 연결 설정
            self.inputs = inputs
            self.outputs = [weakref.ref(output) for output in outputs]
        
        return outputs if len(outputs) > 1 else outputs[0]

 

파이썬에는 with라는 편리한 구문이 있습니다. 아래 코드와 같이 with 블록에 들어갈 때 어떤 처리(전처리)를 해주고, with 블록을 빠져나올 때 처리(후처리)를 자동으로 할 수 있습니다. 해당 코드는 using_config 안에서만 역전파 비활성 모드인 것이고, 빠져나오면 일반 모드인 역전파 활성 모드로 돌아가는 것입니다.

with using_config('enable_backprop', False):
    x = Variable(np.array(2.0))
    y = square(x)

 

앞서 활용했던 using_config 함수는 다음과 같이 구현합니다. 다소 생소할 수 있습니다. 우선 with 블록 안에서 예외가 발생할 것을 고려하여 try/finally 구문을 활용합니다. 미리 이전 값인 old_value를 받아놓고, setattr를 통해 새로운 value를 설정합니다. 최종적으로 with 블록을 빠져나올 때에는 원래 값인 old_value로 복원됩니다.

import contextlib

@contextlib.contextmanager
def using_config(name, value):
    old_value = getattr(Config, name)
    setattr(Config, name, value)
    try:
        yield
    finally:
        setattr(Config, name, old_value)

 

편의성을 위해 다음과 같은 함수를 구현합니다. 이로써 PyTorch의 torch.no_grad와 동일하게 작동시킬 수 있습니다. 앞으로 기울기 계산이 필요 없을 때, no_grad 함수를 호출하면 됩니다.

def no_grad():
    return using_config('enable_backprop', False)
    
with no_grad():
    x = Variable(np.array(2.0))
    y = square(x)