Azure Machine Learning で NumPy の計算速度を高速化する

はじめに

Azure Machine Learning を使うと誰でも簡単に機械学習環境を用意でき、モデル作成や推論用 API をスピーディに公開することができます。ただ、とても簡単に環境を用意できる反面、用意されているライブラリをそのまま学習に用いるとうまく最適化されていなかったりして「うわ、私の Azure ML 遅すぎ…」みたいなことに陥る可能性があります。そうならないように、自分が使用しているライブラリの動作はキチンと知っておこうというのが本記事の趣旨です。

Azure ML で Notebooks を使う

Azure ML のサービスを初めてデプロイして、コンピューティングインスタンスを作成後、ノートブックを開きます。

ノートブック上からはこのような感じでカーネルを選択できると思います。
ターミナルから一覧を取得すると、以下のような出力が得られます。

conda info -e
# conda environments:
#
base                     /anaconda
azureml_py36             /anaconda/envs/azureml_py36
azureml_py38          *  /anaconda/envs/azureml_py38
azureml_py38_pytorch     /anaconda/envs/azureml_py38_pytorch
azureml_py38_tensorflow  /anaconda/envs/azureml_py38_tensorflow

Azure ML の Notebooks にはすぐに利用できるカーネルとして、Python 3.6 と 3.8 系が用意されています。そして Python 3.8 の PyTorch 版と TensorFlow 版も追加で用意されているので、利用時にハードウェアやライブラリ同士の依存性問題に悩むことなくすぐに最新のフレームワークを試せるようになっています。

NumPy の例

Azure ML コンピューティングインスタンス（今回は CPU のみのワークロードを考えます）で思ったよりも速度が出ない一つの例が、NumPy の行列計算です。
NumPy を使って行列の計算をするとき、性能に直結するのが BLAS の実装です。NumPy はバックグラウンドで BLAS という行列演算を行う API 群をコールしますが、BLAS の実装は複数存在（OpenBLAS、ATLAS、Intel Math Kernel Library (Intel MKL) など）していてそれぞれ性能が異なります。

Azure ML の NumPy の BLAS は？

使用している NumPy の BLAS 実装を調べるには、以下のコードを実行します。

import numpy as np 
np.__config__.show()

Azure ML で用意されている 4 つの環境はすべて以下のようになります。

blas_mkl_info:
  NOT AVAILABLE
blis_info:
  NOT AVAILABLE
openblas_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
blas_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
lapack_mkl_info:
  NOT AVAILABLE
openblas_lapack_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
lapack_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
None

blas_opt_info が、['openblas', 'openblas'] となっている場合、OpenBLAS の NumPy ということになります。

BLAS の性能比較

参考までに、Intel 公式による BLAS のベンチマーク結果だと、Intel 製の CPU を使っている場合 ATLAS よりも Intel MKL が高速であることが分かります。

他にも探すといろいろ比較結果が出てきますがおおむね Intel CPU と Intel MKL が一番高速です。

Azure ML で Intel MKL を使う

ここでおすすめなのが、Intel® Distribution for Python* という Intel CPU にゴリゴリに最適化された Python と、NumPy や SciPy などの科学計算ライブラリのパッケージを一括で導入するという案です。

公式のインストール手順はこちらにありますが、こちらのブログで紹介されている超便利な conda 環境構築用 yaml ファイルを参考に、2021 年のリリースノートで Python 3.8 に対応した部分だけ書き換えています。

name: intel-all
channels:
  - intel  #優先的に Intel チャンネルを利用するようにする
  - conda-forge  #Intel チャンネルにないパッケージは conda-forge から入れる
  - nodefaults #default チャンネルを除外する
dependencies:
  - python=3.8 #Intel が配布している pkg は 3.8 まで対応
  - intel::intelpython3_core  #Intel® Distribution for Python を可能な限り使う

以下のコードで、Intel Python 3.8 の環境を作成し、その環境をアクティブにします。

conda env create -f intel-all.yaml
conda activate intel-all

NumPy の BLAS 実装確認コードの実行結果は、以下のようになります。

blas_mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/anaconda/envs/intel-all/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/anaconda/envs/intel-all/include']
blas_opt_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/anaconda/envs/intel-all/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/anaconda/envs/intel-all/include']
lapack_mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/anaconda/envs/intel-all/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/anaconda/envs/intel-all/include']
lapack_opt_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/anaconda/envs/intel-all/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/anaconda/envs/intel-all/include']
None

blas_opt_info の値が ['mkl_rt', 'pthread'] となれば Intel MKL 版が使用されています！

OpenBLAS vs Intel MKL on Azure ML VM

では早速 Azure ML の azureml_py36(OpenBLAS)、azureml_py38(OpenBLAS)、intel-all(Intel MKL) の行列計算ベンチマークを行ってみます。

BLAS を切り替えただけでこんなに性能差が出ました。コードの変更は不要です。Intel MKL いいっ！
他のベンチマークとしては、ベクトル乗算、特異値分解（SVD）、コレスキー分解、固有値分解も試しましたが、総じて Intel MKL が高速です。

比較条件

コンピューティング インスタンス: Standard_F4s_v2 (4 vCPU、8 GB RAM、32 GB ディスク)
CPU: Intel(R) Xeon(R) Platinum 8168 CPU @ 2.70GHz
HyperThreading: ON, 物理コア:2, 論理コア:4
OS: Ubuntu 18.04.6 LTS x86_64
使用コード: 行列乗算用のPythonスクリプト Appendex A: Example
Python: 3.8.11 (default, Sep 9 2021, 00:40:24) [GCC 9.3.0]
numpy: v1.20.3, py38h2742bb4_2, intel
conda 環境: 全リスト

注意点　ハイパースレッディングの機能

CPU のハイパースレッディング機能が ON になっている場合のコアの使い方が、 OpenBLAS と Intel MKL で異なります。BLAS は処理を並列化するのですが、mpstat -P ALL 1 を使用して処理中のコアの使用率を監視すると、Intel MKL は 2 つの論理コアしか使用していないことが分かります。これは Interl MKL が物理コアごとに 1 つの論理コアを使用してパフォーマンスを向上させている仕様とのことです。また、パフォーマンスが低下する恐れがあるので可能であればハイパースレッディングを OFF にすることを推奨しています。

Intel MKL の自動制御を手動で切り、論理コアの最大数まで使用するように変更するには、以下のようにコードに環境変数を追記します。明示的に設定しなければ物理コア数で自動制御されます。

import os
os.environ["MKL_DYNAMIC"] = "FALSE"
os.environ["MKL_NUM_THREADS"] = "4"
# ↑ numpy の import より前に設定する
import numpy as np  

MKL_NUM_THREADS を論理コアの最小～最大に固定してベンチマークを実行すると以下のような結果になりました。

MKL_NUM_THREADS = 4では Intel MKL は行列のサイズによって、遅い/速いが変わってきます。自動制御(2 threads)のままでも十分速いですが、厳密に調整するにはこちらを参考にするとよいかと思います。

他のライブラリ

NumPy のほかにも Intel レポジトリで公開されているライブラリは多岐にわたりますので要チェックです。scypy、scikit-learn、pandas、modin、xgboost、opencv、pytorch、tensorflow など。GPU の場合はまた別の最適化の検討が必要です。

Intel MKL のライセンス

以下、抜粋。

• 使用料はかからない
• パブリックコンテナに含めることができる
• クラウドサーバーにこのリソースをインストールして使用できる
• クラスターにインストールしてユーザーに提供できる
• 再配布は Intel Simplified Software License に従う

まとめ

• Intel MKL いいっ！
• 自分が使ってるライブラリの中身は気にしたほうがいい
• どうやってコンパイルされてるかとか最適化はどうなってるとか
• 環境移行の際の性能問題にも発展しやすいので気を付けよう