Python SVD图像部分重建的快速NumPy表达式

我想获取图像的奇异值分解（分别处理四个RGBa通道中的每一个），然后仅使用第一个

N

向量重建图像。我有下面的代码，正如预期的那样工作，但理想情况下，我希望不使用

np.stack

编写重构。我宁愿使用一些巧妙的广播技巧来达到同样的效果

import matplotlib.image as mpimg
from matplotlib import pyplot as plt

# Load image (python logo from python.org)
img = mpimg.imread('https://www.python.org/static/img/python-logo@2x.png')
# img has shape (164, 580, 4)

# Take SVD, using transpose to get the 4 as
# the first dimension, which makes np.linalg.svd
# treat it as 4 separate 580 x 164 matrices
U, S, V = np.linalg.svd(img.T)
# U has shape (4, 580, 580)
# S has shape (4, 164)
# V has shape (4, 164, 164)

# Reconstruct image with the first N vectors,
# using np.stack(..., axis = -1) to get the 4
# as the last dimension, as in the original image
N = 3
img_reconstructed = np.stack((
    np.dot(U[0,:,:N], S[0,:N,None]*V[0,:N,:]).T, # Red channel
    np.dot(U[1,:,:N], S[1,:N,None]*V[1,:N,:]).T, # Green channel
    np.dot(U[2,:,:N], S[2,:N,None]*V[2,:N,:]).T, # Blue channel
    np.dot(U[3,:,:N], S[3,:N,None]*V[3,:N,:]).T, # alpha channel
), axis = -1)
# img_reconstructed has shape (164, 580, 4), 
# just like the original image

plt.imshow(img_reconstructed)

我想也许我可以用

np.tensordot

实现同样的效果，但我无法让它工作。我试过了

img_reconstructed = np.tensordot(U[:,:,:n], S[:,:n,None]*V[:,:n,:], axes = ((2), (1))).T

但是它的形状是

（164，4580，4）

，而不是我所希望的

（164，580，4）

我设法用

np.einsum

获得了一个有效的解决方案：

img_reconstructed = np.einsum('...ij,...jk->ki...', U[:,:,:n], S[:,:n,None]*V[:,:n,:])

---

但是当我计时时，我发现它比上面使用

np.stack

的解决方案慢25倍。你可以使用

@

而不是

np.dot

进行广播

U, S, V = np.linalg.svd(img.T)
N = 3
img_reconstructed = (U[:,:,:N] @ (S[:,:N,None]*V[:,:N,:])).T

---

matmul/@

更适合批量

dot

当我计时时，我发现它比我第一次尝试使用

np.stack

快五到六倍，所以这正是我想要的。谢谢！