딥러닝 모델을 돌리려면 GPU에 데이터를 올리기위해 DataLoader를 사용하게되는데 DataLoader에서 사용되는 함수가 제너레이터(Generator)이다.

제너레이터는 이터레이터(iterator)를 생성해주는 함수이다. 이터레이터는 __iter__, __next__, __getitem__ 메서드를 구현해야 하지만 제너레이터는 함수 안에서 yield라는 키워드만 사용하면 끝이다. (참고로 제너레이터는 발생자라고 부르기도 한다)

yield란?

그런데 generate라는 키워드를 사용하면 되지 왜 yield라고 이름을 지었을까요? 
yield는 생산하다라는 뜻과 함께 양보하다라는 뜻도 가지고 있습니다. 
즉, yield를 사용하면 값을 함수 바깥으로 전달하면서 코드 실행을 함수 바깥에 양보합니다. 
따라서 yield는 현재 함수를 잠시 중단하고 함수 바깥의 코드가 실행되도록 만듭니다.

함수 안에서 yield를 사용하면 함수는 제너레이터가 되며 yield에는 값(변수)을 지정한다. 아래의 코드에서 next 함수로 __next__ 메서드를 호출하며 yield의 동작 과정을 살펴보자.

def number_generator():
    yield 0    # 0을 함수 바깥으로 전달하면서 코드 실행을 함수 바깥에 양보
    yield 1    # 1을 함수 바깥으로 전달하면서 코드 실행을 함수 바깥에 양보
    yield 2    # 2를 함수 바깥으로 전달하면서 코드 실행을 함수 바깥에 양보
 
g = number_generator()
 
a = next(g)    # yield를 사용하여 함수 바깥으로 전달한 값은 next의 반환값으로 나옴
print(a)       # 0
 
b = next(g)
print(b)       # 1
 
c = next(g)
print(c)       # 2
0
1
2

next(g) 의 반환값을 출력해보면 yield에 지정한 값 0, 1, 2가 차례대로 나온다. 즉, 제너레이터 함수가 실행되는 중간에 next로 값을 가져온다.

동작 과정을 설명하자면,

  1. g = number_generator()와 같이 제너레이터 객체를 만든다.
  2. next(g)를 호출하면 제너레이터 안의 yield 0이 실행되어 숫자 0을 전달한 뒤 바깥의 코드가 실행되도록 양보한다.
  3. 함수 바깥에서는 print(a)로 next(g)에서 반환된 값을 출력

위와 같은 과정을 통해 제너레이터는 함수를 끝내지 않은 상태에서 yield를 사용하여 값을 바깥으로 전달할 수 있다. 즉, return은 반환 즉시 함수가 끝나지만 yield는 잠시 함수 바깥의 코드가 실행되도록 양보하여 값을 가져가게 한 뒤 다시 제너레이터 안의 코드를 계속 실행하는 방식이다.

Generator 사용하기

제너레이터 만들기

range()처럼 동작을 하는 제너레이터를 만들어보자.

def number_generator(stop):
    n = 0              # 숫자는 0부터 시작
    while n < stop:    # 현재 숫자가 반복을 끝낼 숫자보다 작을 때 반복
        yield n        # 현재 숫자를 바깥으로 전달
        n += 1         # 현재 숫자를 증가시킴
 
for i in number_generator(3):
    print(i)

0
1
2

위 코드의 number_generator()함수를 next()를 사용하여 살펴보면 아래와 같다.

>>> g = number_generator(3)
>>> next(g)
0
>>> next(g)
1
>>> next(g)
2
>>> next(g)
Traceback (most recent call last):
  File "<pyshell#100>", line 1, in <module>
    next(g)
StopIteration

값이 3개 밖에 없으므로 next() 를 3번 사용한 후에 다시 호출하면 모든 값을 사용했기 때문에StopIteration 이 발생한다.

yield에서 함수 호출하기

yield에서 함수(메서드)를 호출하면 어떻게 될까? 아래 코드는 리스트에 들어있는 문자열을 대문자로 변환하여 함수 바깥으로 전달한다.

def upper_generator(x):
    for i in x:
        yield i.upper()    # 함수의 반환값을 바깥으로 전달
 
fruits = ['apple', 'pear', 'grape', 'pineapple', 'orange']
for i in upper_generator(fruits):
    print(i)
APPLE
PEAR
GRAPE
PINEAPPLE
ORANGE

리스트 fruits에 들어있는 문자열이 모두 대문자로 출력되었다. 즉, yield에 무엇을 지정하든 결과만 바깥으로 전달한다(함수의 반환값, 식의 결과).

yield from으로 값을 여러 번 바깥으로 전달하기

값을 여러 번 바깥으로 전달할 때는 yield from을 사용하면 된다. yield from에는 반복 가능한 객체, 이터레이터, 제너레이터 객체를 지정한다(yield from은 파이썬 3.3 이상부터 사용 가능).

  • yield from 반복가능한객체
  • yield from 이터레이터
  • yield from 제너레이터객체
def number_generator():
    x = [1, 2, 3]
    yield from x    # 리스트에 들어있는 요소를 한 개씩 바깥으로 전달
 
for i in number_generator():
    print(i)
1
2
3

yield from에 제너레이터 객체 지정하기

def number_generator(stop):
    n = 0
    while n < stop:
        yield n
        n += 1
 
def three_generator():
    yield from number_generator(3)    # 숫자를 세 번 바깥으로 전달
 
for i in three_generator():
    print(i)
0
1
2

number_generator(3)은 숫자를 세 개를 만들어내므로 yield from number_generator(3)은 숫자를 세 번 바깥으로 전달한다. 따라서 for 반복문에 three_generator()를 사용하면 숫자를 세 번 출력한다(next 함수 또는 __next__ 메서드도 세 번 호출 가능).

Reference