ジェネレーター

ジェネレーターはイテレーターの一種である。要素を 1 回取り出す度に何らかの処理を行うときに利用するインターフェースである。例えば、ファイルを処理するとき、10GB にも及ぶ大きなデータをすべてメモリ上に読み込んだ上で処理するよりも、「1 行読んで、処理する」を繰り返した方が、メモリの節約になる。このような順序処理のときにジェネレーターを使用する。

ジェネレーターは関数とほぼ同じように作ることができる。関数では戻り値を返すときは return を使用するが、ジェネレーターでは yield を使用する。yield を使用することにより、関数処理はその行で終了せずに、次の命令を待機するようになる。

例えば、0 から n までの整数列を生成し、その合計を求めるプログラムを考えてみると。整数列を生成するときに、関数を用いると、整数列からなるリストを一端作ってから、そのリストを返す。これに対して、ジェネレーターを用いると、整数を 1 つ生成してはそれを返すという処理を繰り返す。そのため、n が大きくなっても、ジェネレーターでは整数 1 つ分のメモリ量量しか消費しない。

def generate_arr1(n):
    arr = []
    for i in range(n + 1):
        arr.append(i)
    return arr

x = generate_arr1(10)
s = 0
for i in x:
    s = s + i
s
# 55

def generate_arr2(n):
    for i in range(n + 1):
        yield i

x = generate_arr2(10)
s = 0
for i in x:
    s = s + i
s
# 55

ジェネレーターの使用例をもう一つ下に示した。これは、入力されたリストの要素をシャッフルしてから、要素を 2 つずつ返すジェネレーターを作る例である。この種の機能は、機械学習の分野で使われるミニバッチ学習に使われる。例えば、100 万枚の画像を 32 枚ずつ小分けして学習を行う際に 100 万能画像をシャッフルさせて、重複せずにすべて取り出す必要がある。

Python のジェネレーターを使って要素を 2 つずつ取り出す例
import random
random.seed(0)

def get_batch(x):
    random.shuffle(x)
    for i in range(int(len(x) / 2)):
        yield x[(2 * i):(2 * i + 2)]


a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

for b in get_batch(a):
    print(b)

# [7, 8]
# [1, 5]
# [3, 4]
# [2, 0]
# [9, 6]