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

【问题标题】：A better way to do this triple loop using numpy methods使用 numpy 方法执行此三重循环的更好方法
【发布时间】：2017-01-13 21:18:36
【问题描述】：

所以我只是想写一个简单的脚本来将 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 中是否有更好的方法。

【问题讨论】：

是的，因为这就是将 RGB 值转换为 YUV 值的方式
对不起，我忘了补充 X[i, j, k] 是一个 3 维向量（每种颜色一个）
在 NumPy 术语中，这意味着一个 3 elem 数组。因此，相应地进行了编辑。希望没关系。

标签： performance numpy vectorization array-broadcasting

【解决方案1】：

你可以使用np.tensordot -

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

基本上，我们将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

【讨论】：

【解决方案2】：

你可以使用matmul():

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

【讨论】：