计算 3D 和 2D NumPy 数组的加权和

问题描述：

假设我有一个形状为 (n, 3) 的 2D NumPy 数组 A 和一个形状为 (x, y, n) 的 3D 数组 B.我想创建一个形状为 (x, y, 3) 的 3D 数组 Y - 这是一个 RGB 图像.

Let's say I have a 2D NumPy array A of shape (n, 3) and a 3D array B of shape (x, y, n). I want to create a 3D array Y of shape (x, y, 3) - which is an RGB image.

B 的第 3 维包含概率 [0, 1]，其总和为 1.数组 A 是 RGB 颜色列表(每个概率一种颜色).对于 B 的第 3 维中的每个元素 C，我想计算 A * C 的总和，以便获得 RGB 颜色向量.

The 3rd dimension of B contains probabilities [0, 1] which sum up to 1. Array A is a list of RGB colors (one color for each probability). For each element C in the 3rd dimension of B I want to compute the sum of A * C so that I get a RGB color vector.

示例:

给定一个颜色图 A

A = [(255   0   0),
     (  0 255   0),
     (  0   0 255)]

和 B 的第 3 维中点 (x, y) 的元素 C 为

and an element C in the 3rd dimension of B at point (x, y) as

C = [0.2, 0.6, 0.20]

我想将 C' = sum(A * C) 计算为

  (255   0   0) * 0.2
+ (  0 255   0) * 0.6
+ (  0   0 255) * 0.2
---------------------
  ( 51 153  51)

并分配

Y[x, y, :] = C'

我知道我可以用 for 循环遍历 x 和 y 并一次计算每个元素，但我想知道这是否可以以矢量化的方式完成，我不必自己遍历数组(主要是出于性能原因).

I know I could just iterate over x and y with a for loop and compute each element at a time but I wonder if this can be done in a vectorized manner that I don't have to iterate over the array myself (mainly for performance reasons).

答

你是 sum-reducing A 的第一个轴对 B，而其余的轴展开.这是利用基于 BLAS 的张量矩阵乘法的完美设置 - np.tensordot，就像这样 -

You are sum-reducing the first axis from A against the third from B, while the rest of the axes are spread out. This is a perfect setup to leverage BLAS based matrix-multiplication for tensors - np.tensordot, like so -

C = np.tensordot(B,A,axes=((2),(0)))

了解tensordot的相关帖子.

我们也可以手动重塑为2D并对2D使用矩阵乘法:np.dot，像这样 -

We can also manually reshape to 2D and use the matrix-multiplication for 2D : np.dot, like so -

B.reshape(-1,n).dot(A).reshape(x,y,3)

请注意，B.dot(A) 也能正常工作，但速度会较慢，很可能是因为它会循环遍历 B 的第一个轴，同时执行2D 矩阵乘法针对 A 的每个 2D 切片.

Note that B.dot(A) works as well, but that would be slower, most probably as it would loop through the first axis of B, while performing 2D matrix-multiplications for each 2D slice off it against A.

运行时测试 -

In [180]: np.random.seed(0)
     ...: x,y,n = 100,100,100
     ...: A = np.random.rand(n,3)
     ...: B = np.random.rand(x,y,n)

# @Crazy Ivan's soln
In [181]: %timeit np.einsum('ij,kli->klj', A, B)
100 loops, best of 3: 4.21 ms per loop

In [182]: %timeit np.tensordot(B,A,axes=((2),(0)))
1000 loops, best of 3: 1.72 ms per loop

In [183]: np.random.seed(0)
     ...: x,y,n = 200,200,200
     ...: A = np.random.rand(n,3)
     ...: B = np.random.rand(x,y,n)

# @Crazy Ivan's soln
In [184]: %timeit np.einsum('ij,kli->klj', A, B)
10 loops, best of 3: 33.2 ms per loop

In [185]: %timeit np.tensordot(B,A,axes=((2),(0)))
100 loops, best of 3: 15.3 ms per loop

计算 3D 和 2D NumPy 数组的加权和

相关推荐