如何仅用numpy.tensordot实现图像色彩空间转换？

要解决这个问题，核心是理解tensordot的维度收缩规则和输出维度顺序，然后调整参数让它直接输出(m,n,3)的结果，不需要依赖moveaxis。

先理清原有逻辑 #

你原来的高效实现：

np.moveaxis(np.tensordot(transform, image, axes=((-1),(-1))), 0, 2)

是先通过tensordot收缩transform的最后一维（轴1）和image的最后一维（轴2），得到形状为(3, m, n)的中间结果——因为tensordot的输出顺序是第一个数组的未收缩轴在前，第二个数组的未收缩轴在后，这里transform剩下的轴是0（对应3个输出通道），image剩下的轴是0、1（对应图像的高宽），所以输出是(3,m,n)，再通过moveaxis把通道轴移到最后。

直接输出(m,n,3)的tensordot写法 #

我们可以调整输入顺序和收缩的轴，同时转置transform，让tensordot直接输出目标形状：

np.tensordot(image, transform.T, axes=((2), (1)))

为什么这个写法可行？ #

1. 转置transform的作用：原操作是对每个像素pixel执行transform @ pixel（矩阵行与向量点积），转置后的transform.T让我们可以通过向量与矩阵列的点积得到完全相同的结果。
2. 维度收缩规则：我们收缩image的通道轴（轴2）和transform.T的最后一维（轴1），此时：
   • 对于输出的每个(m,n,i)位置，计算的是sum_j image[m,n,j] * transform.T[j,i] = sum_j image[m,n,j] * transform[i,j]，这和transform[i,:] @ pixel（原操作的第i个通道值）完全一致。
3. 输出维度顺序：tensordot会先保留image的未收缩轴（0、1，即高宽），再保留transform.T的未收缩轴（0，即3个输出通道），直接得到(m,n,3)的结果，完美匹配需求。

验证一致性 #

你可以用小数据测试一下：

import numpy as np

# 测试用图像和转换矩阵
image = np.random.rand(2, 2, 3)
transform = np.random.rand(3, 3)

# 原有正确实现
result_einsum = np.einsum('ij,...j', transform, image)
result_moveaxis = np.moveaxis(np.tensordot(transform, image, axes=((-1),(-1))), 0, 2)
# 新写法
result_tensordot = np.tensordot(image, transform.T, axes=((2), (1)))

# 验证结果一致
print(np.allclose(result_einsum, result_tensordot))  # 输出True
print(np.allclose(result_moveaxis, result_tensordot))  # 输出True

内容的提问来源于stack exchange，提问作者emmeowzing