Numpyブロードキャストのテクニックについて

Question

多次元配列のブロードキャストについて質問です。
二次元配列が多次元に並んだ２つのテンソルがあって、そのテンソルの中の行列同士の要素積の最大値を取得し、並べるという操作をしたいと思っています。
for loopで行うとわかりやすいですが計算時間が遅い、ブロードキャストできるようにrepeatで次元を増やすと早くはなるが、メモリを大量に使うためそこがネックに。
もっと賢い方法はありますか？

import numpy as np

X = np.random.randint(0,9,(20,15,13,10,10))
Y = np.random.randint(0,9,(9,10,10))

# X.shape:(A,B,C,D,E)
# Y.shape:(F,D,E)

# for loop
Z = np.zeros((X.shape[0],X.shape[1],X.shape[2],Y.shape[0]))

for a in range(X.shape[0]):
    for b in range(X.shape[1]):
        for c in range(X.shape[2]):
            for f in range(Y.shape[0]):
                Z[a,b,c,f] = np.max(X[a,b,c,:,:]*Y[f,:,:])


# Z.shape:(20, 15, 13, 9)

# ブロードキャスト
X_rep = np.repeat(X[:,:,:,np.newaxis,:,:],1,axis=0)
Y_rep = np.repeat(Y[np.newaxis,np.newaxis,np.newaxis,:,:],1,axis=0)\

Z_rep = np.max(X_rep*Y_rep,axis=(-2,-1))

# Z_rep.shape:(20, 15, 13, 9)

Answer 1

np.max() を計算するときにブロードキャストするのであれば、X の axis=3 にサイズ1の次元を挿入しておくだけで十分だと思います。

ブロードキャスト時の挙動

X: (20, 15, 13, 1, 10, 10)
Y: (9, 10, 10)
↓ 次元を一致させる
X: (20, 15, 13, 1, 10, 10)
Y: ( 1,  1,  1, 9, 10, 10)
↓ 各次元のサイズを一致させる
X: (20, 15, 13, 9, 10, 10)
Y: (20, 15, 13, 9, 10, 10)

import numpy as np

X = np.random.randint(0, 9, (20, 15, 13, 10, 10))
Y = np.random.randint(0, 9, (9, 10, 10))
Z = np.zeros((X.shape[0], X.shape[1], X.shape[2], Y.shape[0]))

for a in range(X.shape[0]):
    for b in range(X.shape[1]):
        for c in range(X.shape[2]):
            for f in range(Y.shape[0]):
                Z[a, b, c, f] = np.max(X[a, b, c, :, :] * Y[f, :, :])

X = np.expand_dims(X, axis=3)
print(X.shape)  # (20, 15, 13, 1, 10, 10)
print(Y.shape)  # (9, 10, 10)
Z2 = np.max(X * Y, axis=(-2, -1))

assert np.allclose(Z, Z2)  # 一致するかチェック

ブロードキャストできるようにrepeatで次元を増やすと早くはなるが、メモリを大量に使うためそこがネックに。

要素数は変わらないので、メモリ使用量は増えないと思います。

追記

質問のループバージョンのコードを numba で最適化したところ、
218ms が 8.45 ms と25倍程度高速化できました。
ブロードキャストで計算したバージョンが 8.01 ms なのでほぼ同じぐらいの速度が出るようになりました。

import timeit
import numpy as np
from numba import jit


np.random.seed(0)
X = np.random.randint(0, 9, (20, 15, 13, 10, 10))
Y = np.random.randint(0, 9, (9, 10, 10))

def calc1(X, Y):
    # broadcast バージョン
    X = np.expand_dims(X, axis=3)
    Z = np.max(X * Y, axis=(-2, -1))


def calc2(X, Y):
    # for-loop バージョン
    Z = np.zeros((X.shape[0], X.shape[1], X.shape[2], Y.shape[0]), dtype=X.dtype)
    for a in range(X.shape[0]):
        for b in range(X.shape[1]):
            for c in range(X.shape[2]):
                for f in range(Y.shape[0]):
                    Z[a, b, c, f] = np.max(X[a, b, c] * Y[f])


@jit(nopython=True)
def calc3(X, Y):
    # for-loop バージョンに numba でコンパイルしたバージョン
    Z = np.zeros((X.shape[0], X.shape[1], X.shape[2], Y.shape[0]), dtype=X.dtype)
    for a in range(X.shape[0]):
        for b in range(X.shape[1]):
            for c in range(X.shape[2]):
                for f in range(Y.shape[0]):
                    Z[a, b, c, f] = np.max(X[a, b, c] * Y[f])

# Jupyter Notebook 上で計測
%timeit calc1(X, Y)
%timeit calc2(X, Y)
%timeit calc3(X, Y)
# 8.01 ms ± 239 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 218 ms ± 32.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# 8.45 ms ± 25.3 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Comments

[todo_suke]

ありがとうございます。
repeatは次元を増やすだけでブロードキャストしてくれたんですね。
確かにメモリ使用量は減らせますが、メモリを一番食うのは
np.max(X_rep*Y_rep,axis=(-2,-1))
の要素積をとった部分なんです。
ここの部分がfor loopだと一個づつ取り出せるのでメモリー消費を抑えることができます。
一つは要素積をとる部分の一部だけにfor loopをかけることかなと思いますが、ほかにアイデアありませんか？

[tiitoi]

> 一つは要素積をとる部分の一部だけにfor loopをかけることかなと思いますが、ほかにアイデアありませんか？

そのバージョンで記述して、numba を使えば高速化できるかもしれません。
https://qiita.com/hanata/items/6003b3b3506e15331ad3

それでも速度がでないようであれば、その部分だけ Cython で書くという手もあります。

[todo_suke]

ありがとうございます。
numba初めて知りました。
コードまで記載していただき助かります。