본문 바로가기
Study 공부할레나

[스터디/밑시딥3] 제 2고지 자연스러운 코드로_18. 메모리 절약 모드

by LENA-cfg 2023. 8. 26.

DeZero의 메모리 사용을 개선할 수 있는 두 가지 구조를 도입합니다.

1. 역전파 시 사용하는 메모리양을 줄이는 방법으로, 불필요한 미분 결고를 보관하지 않고 즉시 삭제

2. 역전파가 필요없는 경우용 모드를 제공(불필요한 계산을 생략)

 

18.1 필요 없는 미분값 삭제

현재 DeZero에서는 모든 변수가 미분값을 변수에 저장하고 있습니다.

y,backward()를 실행하여 미분하면 모든 변수가 미분 결과를 메모리에 유지합니다.

 

머신러닝에서 역전파로 구하고 싶은 미분값은 말단 변수 뿐일때가 대부분이며 중간 변수의 미분값을 필요하지 않습니다.

그래서, 중간 변수에 대해서 미분값을 제거하는 모드를 추가합니다.

class Variable:

    def backward(self, retain_grad=False): # backeard 메서드에 retain_grad를 추가
        if self.grad is None:
            self.grad = np.ones_like(self.data)

        funcs = []
        seen_set = set()

        def add_func(f):
            if f not in seen_set:
                funcs.append(f)
                seen_set.add(f)
                funcs.sort(key=lambda x: x.generation)

        add_func(self.creator)

        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)

Variable 클래스에 backward 메서드에 retaion_grad를 추가합니다.

retain_grad가 True면 모든 변수가 미분 결과(기울기)를 유지합니다.

반대로, False면 중간 변수의 미분값을 모두 None으로 재설정합니다.

 

각 함수의 출력 변수의 미분값을 유지하지 않도록 y().grad = None으로 설정합니다.

이러면 말단 변수 외에는 미분값을 유지하지 않습니다.

18.2 Function 클래스 복습

class Function:
    def __call__(self, *inputs):
        xs = [x.data for x in inputs]
        ys = self.forward(*xs)
        if not isinstance(ys, tuple):
            ys = (ys,)
        outputs = [Variable(as_array(y)) for y in ys]

        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]

 

Dezero에서 미분할때 순전파 => 역전파 순으로 진행되고 역전파 시 순전파의 계산 결과가 필요해서 결과 값을 기억해둡니다. 해당 과정에서 순전파의 결과값을 보관하는 부분이 self.inputs = inputs입니다.

 

함수는 입력을 inputs이라는 '인스턴스 변수'로 참조합니다.

그 결과 inputs이 참조하는 변수의 참조 카운트가 1만큼 증가하고 __call__메서드에서 벗어난 뒤에도 메모리에 생존합니다. 인스턴스 변수인 inputs로 참조하지 않았다면 참조 카운트가 0이되어 메모리에서 삭제됐을 겁니다.

 

미분값이 필요없는 경우엔 중간 계산 결과를 저장, 연결한 필요가 없습니다.

* 신경망에는 학습과 추론이라는 두 단계가 있는데 학습 때에는 미분값을 구해야하지만,

추론 시에는 단순히 순전파만 하기 때문에 중간 계산 결과를 곧바로 버리면 메모리 사용량을 줄일 수 있습니다.

18.3 config 클래스를 활용한 모드 전환

'역전파 활성 모드'와 '역전파 비활성 모드'를 전환하는 구조가 필요합니다.

Config 클래스를 이용합니다.

class Config:
	enable_backprop = True

해당 클래스는 불리언 타입의 enable_backprop만 존재하며 역전파 가능 여부를 뜻하고 

True일때 역전파 활성 모드가 됩니다.

class Function:
    def __call__(self, *inputs):
        xs = [x.data for x in inputs]
        ys = self.forward(*xs)
        if not isinstance(ys, tuple):
            ys = (ys,)
        outputs = [Variable(as_array(y)) for y in ys]

        if Config.enable_backprop:
            self.generation = max([x.generation for x in inputs]) # 1. 세대 설정
            for output in outputs:
                output.set_creator(self) # 2. 연결 설정
            self.inputs = inputs # 순전파의 결과값을 보관
            self.outputs = [weakref.ref(output) for output in outputs]

        return outputs if len(outputs) > 1 else outputs[0]

if Config.enable_backprop:

이 코드는 Config.enable_backprop이 True일 때만 역전파 코드가 실행됩니다.

 

1. 세대 설정 에서 정하는 세대는 역전파 시 노드를 따라가는 순서를 정하는 데 사용됩니다.

즉, '역전파 비활성 모드'에서는 필요하지 않습니다.

2. 연결 설정 부분은 계산들의 '연결'을 만드는데, '역전파 비활성 모드'에서는 필요 없습니다.

18.4 모드 전환

Config.enable_backprop = True
x = Variable(np.ones((100, 100, 100)))
y = square(square(square(x)))
y.backward()
print(x.grad)

Config.enable_backprop = False
x = Variable(np.ones((100, 100, 100)))
y = square(square(square(x)))
print(x.grad)

enable_backprop = True 면 중간 계산 결과가 역전파 완료 되기 전까지 메모리를 차지합니다.

enable_backprop = False 면 중간 계산 결과는 사용 후 곧바로 삭제됩니다.(다른 객체에서 참조가 없어지는 시점에 메모리에서 삭제됩니다.)

18.5 with 문을 활용한 모드 전환

후처리를 자동으로 수행할때 사용하는 구문인 with를 사용합니다.

예시로 파일의 open과 close입니다.

f = open('samle.txt', 'w')
f.write('hello world!')
f.close()

위의 코드처럼 매번 close()하는게 귀찮으니 with문을 사용해 귀찮음을 줄여줍니다.

with open('sample.txt', 'w') as f:
	f.write('hello world!')

with 안으로 들어갈 때의 처리(전처리)와 with블록을 빠져나올 떄의 처리(후처리)를 자동으로 할 수 있습니다.

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

with using_config('enable_backprop', False): 안에서만 '역전파 비활성 모드'가 됩니다. 

그리고 with 블록을 벗어나면 일반 모드, 즉 '역전파 활성 모드'가 됩니다.

import contextlib

@contextlib.contextmanager
def config_test():
    print('start')
    try:
        yield
    finally:
        print('done')
        
with config_test():
    print('process...')

@contextlib.contextmanager 데코레이터를 달면 문맥을 판단하는 함수가 만들어지고

함수 안에서 yield 전에는 전처리 후에는 후처리 로직이 됩니다.

 

with config_test():
    print('process...')

해당 구문을 사용하면 with 블록 안으로 들어갈 때 전처리가 실행되고 블록 범위를 빠져나올 대 후처리가 실행됩니다.

* with 안에서 예외가 발생할 수 있고 예외는 yield를 실행하는 코드로 전달되어 try/finally로 감싸 처리합니다.

import contextlib

class Config:
    enable_backprop = True

@contextlib.contextmanager
def using_config(name, value): # name(str)은 Config 속성의 이름(클래스 속성 이름)을 가리킴
    old_value = getattr(Config, name) # getattr함수에 넘겨 config 클래스에서 꺼냄
    setattr(Config, name, value) # setattr 함수를 사용하여 새로운 값을 설정
    try:
        yield
    finally:
        setattr(Config, name, old_value) 
        # with 블록으로 들어갈 때 name으로 지정한 config 클래스 속성이 value로 설정되고
        # with 블록을 빠져나오면서 원래 값(old_value)로 복원됨
        
# 실제 using config 함수 사용        
with using_config('enable_backprop', False): 
    x = Variable(np.array(2.0))
    y = square(x)

with using_config('enable_backprop', False):를 매번 적기는 귀찮아서 no_grad라는 함수만들어줍니다.

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