如何用np.tensordot替代循环实现张量与另一张量连续切片相乘？

问题描述 #

现有两个3D张量，形状分别为 T1=(a,b,c) 和 T2=(d,b,c)，需要将T2分别与T1的每一个(b,c)维度切片进行元素级相乘，最终得到形状为(a,d,b,c)的结果。已有迭代实现代码如下：

import numpy as np

a=2
b=3
c=4
d=5

T1 = np.random.rand(a,b,c)
T2 = np.random.rand(d,b,c)

# 迭代实现
L= []
for j in range(a) :
    L+=[T1[j,:,:]*T2]
L = np.array(L)
print(L.shape)  # 输出 (2,5,3,4)

寻求非迭代的numpy实现方法，避免使用循环。

解决方案 #

利用numpy的广播机制可以直接实现该操作，无需循环。核心思路是给两个张量分别添加一个新维度，让它们的维度对齐后自动广播相乘：

方法1：直接扩展维度实现广播 #

通过在T1的第1位添加维度，T2的第0位添加维度，使两者的维度变为(a,1,b,c)和(1,d,b,c)，此时numpy会自动将这两个维度广播为(a,d,b,c)，再进行元素级相乘：

# 非迭代实现
result = T1[:, None, :, :] * T2[None, :, :, :]
print(result.shape)  # 输出 (2,5,3,4)

方法2：使用np.expand_dims显式扩展维度 #

如果觉得直接索引添加维度不够直观，可以用np.expand_dims显式指定扩展位置：

T1_expanded = np.expand_dims(T1, axis=1)  # 形状变为 (a,1,b,c)
T2_expanded = np.expand_dims(T2, axis=0)  # 形状变为 (1,d,b,c)
result = T1_expanded * T2_expanded
print(result.shape)  # 输出 (2,5,3,4)

验证结果 #

可以通过np.allclose验证非迭代方法和迭代方法的结果是否完全一致：

print(np.allclose(result, L))  # 输出 True

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