numpyの内積計算

Question

2次元データのaとbとの内積結果をcに保存したいときに,
1次元のnp.dotを繰り返して計算しなければならず、np.dotを2次元に拡張して計算することはできますか？
下for文のでの実装をしましたが、計算量が多く時間がかかるため高速化を行いたいです。

import numpy as np

a = np.arange(100 * 8*20*20).reshape(100,8*20*20)
b = np.arange(100 * 8).reshape(100,8)
c = np.zeros(100)
for i in range(100):
    for j in range(400):
        c[i] += np.dot(a[i,8*j:8*(j+1)],b[i,:])
print(c)

Comments

[tiitoi]

質問のコードを実行すると reshape のエラーがでます。配列の形状を確認してください。

[redcatML]

プログラム修正しました。
確認いただけたらと思います

Answer 1

一行で書くなら以下のようになります。

c = (a.reshape(100, -1, 8) * b.reshape(100, 1, 8)).sum(axis=(1, 2))

import numpy as np

a = np.arange(100 * 8 * 20 * 20).reshape(100, 8 * 20 * 20)
b = np.arange(100 * 8).reshape(100, 8)

## 質問のコード
c = np.zeros(100)
for i in range(100):
    for j in range(400):
        c[i] += np.dot(a[i, 8 * j : 8 * (j + 1)], b[i])

## 回答のコード
c2 = (a.reshape(100, -1, 8) * b.reshape(100, 1, 8)).sum(axis=(1, 2))

# 結果が一致するかどうか
assert np.allclose(c, c2)

timeit のベンチマークによると、ループを削除したことにより 300倍ぐらい高速になりました。

## 質問のコード
143 ms ± 2.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
## 回答のコード
553 µs ± 384 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Comments

[redcatML]

回答ありがとうございます。
コード確認しました。
reshapeとsumの組み合わせをよく考えてみます

[tiitoi]

コードに解説を書きませんでしたが、
a.shape: (100, 3200), b.shape: (100, 8)
なので、
a.shape: (100, 400, 8), b.shape: (100, 1, 8)
に reshape してから乗算すると、b は (100, 400, 8) にブロードキャストされて (axis=1 の値は 400個同じ値が繰り返される) から、計算されます。
その結果、(100, 400, 8) という乗算結果が出てくるので、axis=(1, 2) の sum() をとれば内積を計算できたことになり、最終的に c.shape: (100,) が出てきます

numpy のブロードキャストの仕様についてわからなければ、そこから勉強してみるといいかと思います。