[Python] 예제를 통해 알아보는 파이썬 yield 활용법, 메모리 사용량 비교하기

yield 사용법

yield는 제네레이터를 정의하기 위해 function에서 사용되는 파이썬의 키워드이다.

참고로 제너레이터는 실행을 잠깐 멈추고 결과를 리턴한 다음 다시 제너레이터를 호출할 때 멈췄던 위치에서 실행을 재개할 수 있는 특수한 종류의 함수이다.

즉, 값을 한꺼번에 생성하여 메모리에 저장할 필요가 없이 필요할 순간마다 값을 생성할 수 있다.

 

다음은 피보나치 수열을 생성할 때 yield을 사용하는 제너레이터 함수의 예이다.

def fibonacci():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

이 코드를 실행시켜보자.

fib = fibonacci()
for i in range(10):
    print(next(fib))

그럼 다음과 같은 결과를 얻는다.

0
1
1
2
3
5
8
13
21
34

위 예는 설명을 위한 간단한 코드라서 메모리를 많이 차지하지 않는다.

yield가 유용한 한 가지 상황은 한 번에 어떤 시퀀스를 생성할 때 많은 메모리를 소모할 경우이다.

yield를 사용한 제너레이터 함수를 만들면 각 값을 필요할 때마다 생성해서 시퀀스 전체를 메모리에 모두 저장할 필요가 없어진다.

---

메모리 사용량 비교

yield을 살펴보는 김에 메모리 사용량을 비교해보자.

먼저 같은 역할을 하는 코드이지만 yield을 사용하는 버전과 아닌 버전을 살펴보자.

def fibonacci_list(n):
    fibonacci = [0, 1]
    for i in range(2, n):
        fibonacci.append(fibonacci[-1] + fibonacci[-2])
    return fibonacci[:n]

def fibonacci_generator(n):
    a, b = 0, 1
    for i in range(n):
        yield a
        a, b = b, a + b

두 코드 모두 피보나치 수열을 출력하는 코드이다.

먼저 나온 코드는 yield을 사용하지 않은 버전이고 아래 코드는 사용한 버전이다.

두 코드를 비교할 때 비슷한 환경을 만들어주기 위해 yield을 사용한 코드를 본문에 있는 코드에서 약간 바꾸었다.

 

메모리 사용을 비교하는 코드를 완성시켜보자.

import psutil

def fibonacci_list(n):
    fibonacci = [0, 1]
    for i in range(2, n):
        fibonacci.append(fibonacci[-1] + fibonacci[-2])
    return fibonacci[:n]

def fibonacci_generator(n):
    a, b = 0, 1
    for i in range(n):
        yield a
        a, b = b, a + b

if __name__ == '__main__':
    n = 100000
    mem_before_list = psutil.Process().memory_info().rss / 1024 / 1024
    fibonacci_list(n)
    mem_after_list = psutil.Process().memory_info().rss / 1024 / 1024
    print(f"Memory used by list function: {mem_after_list - mem_before_list} MB")

    mem_before_gen = psutil.Process().memory_info().rss / 1024 / 1024
    gen = fibonacci_generator(n)
    for i in gen:
        pass
    mem_after_gen = psutil.Process().memory_info().rss / 1024 / 1024
    print(f"Memory used by generator function: {mem_after_gen - mem_before_gen} MB")

psutil을 사용하여 간단하게 구현할 수 있다.

 

이 코드를 돌린 결과는 다음과 같다.

n=100000일 때, 결과 비교

파일 이름을 test.py로 저장하고 코드를 실행했을 때 yield을 사용하면 메모리를 거의 사용하지 않은 것을 확인할 수 있다.

n 값을 200000으로 키운 다음 다시 돌려보자.

n=200000일 때, 결과 비교

결과에서 볼 수 있는 것처럼 코드에서 n 값이 커질수록 리스트 방식에서는 메모리 사용량이 늘어나지만 yield에서는 큰 변화가 없다.

---

살펴본 것처럼 yield을 사용하면 프로그램을 보다 가볍게 만들 수 있으니 적재적소에 잘 활용해보자.