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]))

基本上，我们将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方法执行此三重循环的更好方法

2 个答案:

基准