【Python】Generatorを使うメリット

この記事では、PythonのGenerator(ジェネレータ)を使うメリットをご紹介します。

ジェネレータはリストに似ていますが、メモリ消費の観点からみると全然違います。

Listを使った処理ではメモリを消費しますが、generatorを使うとメモリを消費しません。

その違いについてコードを見ながら説明していきます。

リスト

リストで与えた数値の三乗を計算して返す関数を作ってみます。

def get_cube(numbers):
  r = []
  for n in numbers:
    r.append(n*n*n)
  return r

numbers = [1,3,5,7,9,11]

cubes = get_cube(numbers)

print(type(cubes))

print(cubes)

実行結果は以下の通り。

<class 'list'>

[1, 27, 125, 343, 729, 1331]

リストとして計算結果が返ってきました。

ジェネレータ

ジェネレータを使って、同じ関数を書いてみます。

ジェネレータを使うときには、returnの代わりにyieldを使います。

def cube_generator(numbers):
  for n in numbers:
    yield n*n*n

gen = cube_generator(numbers)

print(type(gen))

for i in gen:
  print(i)

実行結果は以下の通り。

<class 'generator'>
1
27
125
343
729
1331

全く同じ結果が得られました。

Generatorクラスは、リストと同じようにfor文で回して値を取り出すことができます。

メモリ使用量の違い

リストでもジェネレータでも同じ結果が得られました。

それではジェネレータを使うメリットとは何でしょうか?

メモリ使用量を見てみましょう。Memory_profilerを使うと良いです。

Google colabではデフォルトで入っていなかったので、pip installでインストールしました。

import memory_profiler

まずはリストの場合の実行結果です。

今度はわかりやすいように1~100万までの数値の3乗を返すようにします。

リストの数値が少ないとあまりメモリを消費しませんからね。

メモリ使用量だけでなく、timeで時間も計測します。

%time

m1 = memory_profiler.memory_usage()

cubes = get_cube(range(1000000))

m2 = memory_profiler.memory_usage()

print("memory usage[M]", m2[0]-m1[0])
CPU times: user 4 µs, sys: 1e+03 ns, total: 5 µs

Wall time: 8.11 µs

memory usage[M] 55.1953125

8マイクロ秒で、55Mのメモリを消費したことが分かりました。

同様のコードをジェネレータで実装してみます。

%time

m1 = memory_profiler.memory_usage()

gen = cube_generator(range(1000000))

m2 = memory_profiler.memory_usage()

print("memory usage[M]", m2[0]-m1[0])
CPU times: user 2 µs, sys: 0 ns, total: 2 µs

Wall time: 7.39 µs

memory usage[M] 0.0

実行時間はほとんど同じで、メモリ使用量を0Mに抑えられました。

全く同じ処理をしているのですが、メモリを消費せずに済みます

しかも、実行速度はほとんど変わりません。

これがジェネレータを使う大きなメリットです。

参考

https://medium.com/@chetaniam/optimize-python-code-with-generators-aef839996ee4