BETA

【NumPy】2つの行列のすべての行同士で差を取るときのベストプラクティス

投稿日:2020-01-16
最終更新:2020-01-16

2つの行列$A$、$B$とあり、$A$の各行と$B$の各行の差の計算をすべての組み合わせでおこないたいことが稀にあります(だいたいRBFカーネルの計算です)。
今回は思いつく3つの方法で速度を比較してみました。

準備

以下で使われる$A$と$B$は次のように生成しているとします。

import numpy as np  

A = np.random.rand(n, dim)  
B = np.random.rand(n, dim)  

上記のnとdimですが、次の4つの組み合わせで実験をおこなっています。

  • n=100, dim=100
  • n=1000, dim=100
  • n=100, dim=1000
  • n=1000, dim=1000

また、計算は4コアのMac Book Proを使っておこなっています。

1. 愚直に2重forループ

愚直に書くと次のようにforループを2つ使って実現できます。

diffs = np.empty([n, n, dim])  
for i, a in enumerate(A):  
    for j, b in enumerate(B):  
        diffs[i, j, :] = a - b  

このときの速度は以下のようになりました。

  dim=100 dim=1000
n=100 12.5ms 30.9ms
n=1000 1.52s 6.38s

2. ブロードキャストを利用

よく知られているようにforを繰り返すと計算効率が悪くなります。
NumPyのブロードキャストを利用することで、forループを1回減らすことができます。

diffs = np.empty([n, n, dim])  
for i, a in enumerate(A):  
    diffs[i, :, :] = a - B  

このときの速度は以下のようになりました。

  dim=100 dim=1000
n=100 2.84ms 18.6ms
n=1000 451ms 6.04s

ブロードキャストを利用することで全体的に計算速度が向上しました。
特にn=100, dim=100のときには4倍以上速くなっています。

3. forループを使わない

forループを1回減らすことができるなら、0回にもできるんじゃないかと思いますよね。
実は次のようにブロードキャストを上手く利用することで実現できます。

diffs = A.reshape(n, 1, dim) - B.reshape(1, n, dim)  

reshapeすることで3階のテンソルになっているので複雑なように思えます。しかし、最後のdimのサイズのところを無視して考えてみると、$n$次元の列ベクトルと$n$次元の行ベクトルでのブロードキャストと一緒です。

このときの速度は以下のようになりました。

  dim=100 dim=1000
n=100 1.75ms 17.9ms
n=1000 405ms 4.61s

計算速度のまとめ

ここまで紹介した方法との比較は以下のとおりです。

  n=100, dim=100 n=100, dim=1000 n=1000, dim=100 n=1000, dim=1000
1の方法 12.5ms 30.9ms 1.52s 6.38s
2の方法 2.84ms 18.6ms 451ms 6.04s
3の方法 1.75ms 17.9ms 405ms 4.61s

どのケースにおいても3の方法が一番速いことがわかります。

もっと速い方法求む

タイトルにはベストプラクティスとあるのですが、もっと速い方法ないのかなぁと思っておりますので我こそはという方はぜひ教えて下さい!

技術ブログをはじめよう Qrunch(クランチ)は、プログラマの技術アプトプットに特化したブログサービスです
駆け出しエンジニアからエキスパートまで全ての方々のアウトプットを歓迎しております!
or 外部アカウントで 登録 / ログイン する
クランチについてもっと詳しく

この記事が掲載されているブログ

@opqrstuvcutの技術とか技術以外のブログ 最近フリーランスになりました

よく一緒に読まれる記事

0件のコメント

ブログ開設 or ログイン してコメントを送ってみよう