使用 numpy 方法执行此三重循环的更好方法

Question

所以我只是想写一个简单的脚本来将 RBG 转换为 YUV，结果是这样的：

rgb2yuv_matrix = np.array([[0.299, 0.587, 0.114], [-0.1473, -0.28886, 0.436],[0.615, -0.51499, 0.10001]])
for i in range(n_train):
    for j in range(32):
        for k in range(32):
            # Note here that X_train[i, j, k] is a 3 element array
            X_train[i,j,k] = np.dot(rgb2yuv_matrix, X_train[i,j,k])

我想知道是否有更好的方法来做到这一点，因为我知道三重循环非常低效，但我不知道在 numpy 中是否有更好的方法来做到这一点。

Answer 1

可以使用np.tensordot-

np.tensordot(X_train, rgb2yuv_matrix, axes=([-1],[1]))

基本上，我们是 sum-reducing X_train 的最后一个轴相对于 rgb2yuv_matrix 的第二个（也是最后一个）轴。

对于具有 np.dot 的小型数组，您可能会获得一些性能，这也涉及一些重塑，就像这样 -

X_train.reshape(-1,3).dot(rgb2yuv_matrix.T).reshape(X_train.shape)

基准测试

In [25]: def original_app(X_train, rgb2yuv_matrix, n_train, m):
    ...:     for i in range(n_train):
    ...:         for j in range(m):
    ...:             for k in range(m):
    ...:                 X_train[i,j,k] = np.dot(rgb2yuv_matrix, X_train[i,j,k])
    ...:                 

In [26]: # Setup inputs
    ...: n_train = 1000
    ...: m = 32
    ...: X_train = np.random.rand(n_train,m,m,3)
    ...: rgb2yuv_matrix = np.array([[0.299, 0.587, 0.114], \
    ...:                            [-0.1473, -0.28886, 0.436],\
    ...:                            [0.615, -0.51499, 0.10001]])
    ...: 

In [27]: %timeit original_app(X_train, rgb2yuv_matrix, n_train, m)
1 loops, best of 3: 2.45 s per loop

In [28]: %timeit np.tensordot(X_train, rgb2yuv_matrix, axes=([-1],[1]))
10 loops, best of 3: 22 ms per loop

In [29]: %timeit X_train.reshape(-1,3).dot(rgb2yuv_matrix.T).reshape(X_train.shape)
10 loops, best of 3: 22.6 ms per loop

Answer 2

您可以使用 matmul():

a = np.random.rand(3, 3)
b = np.random.rand(10, 32, 32, 3)
c = np.matmul(a, b[..., None]).squeeze()

使用 numpy 方法执行此三重循环的更好方法

A better way to do this triple loop using numpy methods

performance

numpy

vectorization

array-broadcasting

基准测试