Numpy/PyTorchのブロードキャスト機能の解説

タグ：python machine_learning

行列演算ライブラリのブロードキャスト機能は、形状の違う行列間の演算を可能にする機能である。数学上はこのようなときにブロードキャストは小さいほうの行列の形状をうまい具合に調整してくれる。ところがこの強力な機能には地味な制約がありバグのもとになるので仕様をよく理解して活用したい。

基本的な例

ブロードキャストがうまく行くための条件は、次元を末尾からチェックしていって

双方の行列の末尾側次元が演算可能な大きさになっている
先頭側の演算可能でなかった次元は、どちらかが1であるか、存在しない

となっていることである。わかりにくいので図にすると

のようなチェックが走っていることになる。

実行して確認してみると

    In [1]: a = np.ones((2, 3, 4, 5))

    In [2]: b = np.ones((1, 4, 5))

    In [3]: (a + b).shape
    Out[4]: (2, 3, 4, 5)

のように次元が大きいほうの行列にそろえられた。
また大きさ1の次元は自動でつくられるのでなくてもよい

    In [1]: a = np.ones((2, 3, 4, 5))

    In [2]: b = np.ones((4, 5))

    In [3]: (a + b).shape
    Out[4]: (2, 3, 4, 5)  # 同じ

応用編：ちょっとわかりにくいがうまくいく例

大きさ1の次元が演算対象の行列双方にあっても問題ない。このときは自動で1出ないほうにそろえられる。

実行例：

    In [1]: a = np.ones((2, 1, 4, 5))

    In [2]: b = np.ones((1, 3, 4, 5))

    In [3]: (a + b).shape
    Out[4]: (2, 3, 4, 5)  # 同じ

ちょっとわかりにくいが、行列積 (@演算子) のように異なる大きさの行列を対象とする演算もうまくいく

    In [1]: a = np.ones((2, 3, 4, 5))

    In [2]:b = np.ones((1, 5, 8))

    In [3]: (a @ b).shape
    Out[4]: (2, 3, 4, 8)

この場合は (4, 5) 行列と (5, 8) 行列の行列積が実行され、 (4. 8) 行列 (をさらに2*3個並べたもの) ができた。

できない例

末尾側の次元に1を入れればそろう (2×3×2 vs 2×3 など) ときでも末尾側への次元挿入はできない
要素数1の次元を削除すればそろう例 (2×1×3 vs 2×3 など) でも削除はできない

このような場合はunsqueeze/squeezeなどで手でそろえてから演算する必要がある。

Numpy/PyTorchのブロードキャスト機能の解説

基本的な例

応用編：ちょっとわかりにくいがうまくいく例

できない例

おすすめ記事