NumPy をマルチスレッドで計算させる

Python

Anaconda / Miniconda や pip で普通に NumPy をインストールすると、その計算は複数 CPU コアを備えた計算機でも常にシングルスレッドで実行され…ると今日まで思い込んでいた。が、実は環境変数を設定すれば NumPy はマルチスレッドで計算してくれるらしい。もっと早く知っておきたかったな。。。

ということで、今日は NumPy のマルチスレッド計算を有効化する環境変数について備忘録。

スレッド数を指定する環境変数

変数名

NumPy の計算のうち、コアの部分は BLAS (Basic Linear Algebra Subprograms) ライブラリで実行される。この BLAS ライブラリ自身が、特定の環境変数で指定されたスレッド数で計算するように作られている。なので NumPy が使用中の BLAS ライブラリ用の環境変数に、計算に使いたいスレッド数を指定すれば良い。

本日 2020-09-20 時点で調べられた範囲で書いておくと、BLAS ライブラリごとのスレッド数を指定する環境変数は次の表のようになっていた:

BLAS ライブラリ 環境変数
OpenBLAS OPENBLAS_NUM_THREADS
MKL (Intel CPU 用) MKL_NUM_THREADS
veclib (macOS) VECLIB_MAXIMUM_THREADS
ATLAS 無し

ATLAS は Automatically Tuned Linear Algebra Software という名前の通り、計算に使うスレッド数も自動的に決定するらしく、環境変数での設定はできないらしい。

設定方法

普通に export MKL_NUM_THREADS=4 (Windows ならば SET MKL_NUM_THREADS=4) などと環境変数を設定すれば良い。また、これらの環境変数は NumPy が import されるまでに設定されていれば良いため、次のように書いても問題無い（この方法は Jupyter Notebook などで便利だ）:

import os
os.environ["OPENBLAS_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
os.environ["VECLIB_NUM_THREADS"] = "4"

import numpy as np

# ...以下略...

なお NumPy (MKL?) は最後に使った MKL_NUM_THREADS の値を覚えているのか、当該環境変数を設定して実行した後、当該環境変数を未定義状態にして実行しても前回と同じ数のスレッドを計算に使うようだった。このあたりは調べておらず詳細不明だけれど、期待していない挙動をしたら明示的にスレッド数を指定した環境変数定義を行った方が無難であるもよう。

注意書き

ヘヴィな科学計算を含んだプログラムであれば multiprocessing や joblib 、あるいは Cython+OpenMP や Numba の prange などを使って並列演算を自分で書いたりすることもあると思う。こういうケースでは、最終的に「枝葉」の処理で実行される NumPy の計算が何スレッドで実行されるのかを把握しておかないとスレッド数過多になって逆に遅くなる可能性がある。こういう話をキチンと考えていくのは泥沼的に面倒くさいので、こういう面倒事が嫌なら NumPy (BLAS) レベルでの並列計算をあえて行わない、という選択肢もアリかなと思う。

NumPy が使用している BLAS の確認方法

NumPy が使用している BLAS は numpy.__config__.blas_opt_info という辞書オブジェクトの libraries キーで確認できるようだ。

たとえば Core i7 を搭載した手元の計算機 (Ubuntu) にて conda install numpy した場合、次のように mkl_rt と pthread という値が得られた:

$ python -c "import numpy; print(numpy.__config__.blas_opt_info['libraries'])"
['mkl_rt', 'pthread']

おそらく、MKL を使っているのだろう…ということで MKL_NUM_THREADS を設定してみると、確かに指定した数のスレッドで計算が実行された。

今度は、同じ計算機で pip install numpy で NumPy をインストールしたところ、次のように openblas という値が得られた:

$ python -c "import numpy; print(numpy.__config__.blas_opt_info['libraries'])"
['openblas', 'openblas']

なるほど OpenBLAS だろうかと思って OPENBLAS_NUM_THREADS を設定したところ、確かに指定した数のスレッドで計算が実行された。

以上