[python] Numpyで巨大な配列を効率的に処理するためのTips集

機械学習などで大規模なデータを扱っていると、ごくごく自然にそれらのデータをpandasやnumpy配列に突っ込んで処理したくなるときがあります。
Pythonの場合、forループを回すよりもnumpyの関数を使ったほうが圧倒的に高速なケースは多々あり、日々numpyの恩恵を受けています。

それでも、例えば数万✕数万の行列の演算をしようとすると途方もない時間がかかる、もしくはデータが巨大すぎて１つの配列にデータを保持したまま演算しようとするとメモリ不足でプロセスが強制終了するなど悩まされるケースもあります。
わたし自身、最近この問題で非常に悩まされたので、同じくビッグデータを扱うPythonユーザーたちの助けになればとこの記事を書きました。

Numpyでも遅い問題

Numpyは普通に演算する分にはとても早いです。しかしそれでも配列のサイズが大きくなると演算が遅くなるケースがあります。その例をまずは見てみましょう。

使用したPCのスペックは以下の通りです。

マシン	MacBook Pro (13-inch)
CPU	2.5 GHz Intel Core i7
メモリ	16 GB
コアの個数	2

例として、10,000 ✕ 10,000の2次元配列の内積計算をしてみます。

import numpy as np
import time

# np.ones() で取得できるデータ型は float64
arr = np.ones((10000,10000))

def npdot_time(arr):
  start = time.time()
  np.dot(arr, arr)
  return time.time() - start
  
npdot_time(arr)
# 30.2010498046875

10,000行なんてビッグデータでも何でもないような行数だと思われますが、実際に演算は30秒ほどかかっています。桁が増えるとどれほどの処理時間になるか想像したくないほどです。

今回はこの行列の内積演算を高速化・効率化してみましょう。

高速化・効率化Tips

1. データ型をfloat32にする

np.float32 は32ビット、つまり単精度の浮動小数点数を表しています。このデータ型を指定することでメモリを効率的に使うことができます。

実際に float32を指定した場合の速度を見てみましょう。

arr = np.ones((10000,10000), dtype='float32')

npdot_time(arr)
# 14.641755819320679

データ型を指定しただけで実行時間が半分以下の14秒台まで出ました。これで倍速はなかなか強いです。

2. daskを使う

daskとは、numpyやpandasライクで並列処理、分散処理を行うことのできるライブラリです。入れてない人は

pip install dask

で入れておいてください。
daskの使い方は以下の記事が詳しいので参照してみてください。
データ分析のための並列処理ライブラリDask

では実際にdaskで内積演算を書いてみましょう。

import dask.array as da

arr = np.ones((10000,10000))
# chunks は配列の長さの半分の 5000 の指定とする
darr = da.from_array(arr, chunks=(5000, 5000))

def dadot_time(arr):
  start = time.time()
  dadot = da.dot(darr, darr)
  dadot.compute() 
  return time.time() - start
  
dadot_time(darr)
# 28.194761991500854

30秒かかっていた処理からは2秒短縮です。このくらいの配列のサイズだと効果を実感しにくいですが、サイズが大きくなるに連れて分散処理の恩恵を感じられるようになります。

3. 行列を分割して処理をする

そもそも、配列が大きすぎるので処理に時間がかかったり、メモリに乗らなかったりするわけです。であれば、配列自体を小さくして処理するのが妥当でしょう。

今回は行数を100分割して処理をしてみようと思います。

arr = np.ones((10000,10000))
arrs = np.array_split(arr, 100)

def npdot_split_time(arrs, arr):
  start_time = time.time()
  dots = [np.dot(a, arr) for a in arrs]
  dot_matrix = np.concatenate(dots)
  return time.time() - start_time
  
npdot_split_time(arrs, arr)
# 27.128658771514893

これも30秒かかっていたころから3秒の短縮です。これも速度を劇的に改善する手法ではないのですが、配列のサイズが大きすぎて内積計算中にプロセスが強制終了（KILL 9）されることを防ぐための方法となります。

まとめ

各手法の実行時間とどんなときに使えるかをまとめました。実際にはこれらを合わせて使うことで巨大な配列の演算を実現させます。

手法	実行時間[s]	どんなときに使えるか
そのまま numpy.dot	30.2	---
float32に変更	14.6	単純に関数の処理を高速化したい
daskを使う	28.2	分散処理を使って複数のCPUで処理したい
配列を分割する	27.1	メモリに乗らないくらいの巨大な配列計算を処理したい

これらのTipsが皆さんの参考になれば幸いです！