TensorFlow中特定张量逐元素乘法：如何避免使用for循环

方法1：利用广播与求和简化（最高效） #

你的操作本质是每个空间位置(i,j)处，将a[i,j,:]的3个通道求和后，与b[i,j,f]逐元素相乘（通过广播扩展维度）。直接用以下代码即可完成：

import tensorflow as tf

# 假设a的形状为(224, 224, 3)，b的形状为(224, 224, f)
sum_a = tf.reduce_sum(a, axis=-1)  # 形状变为(224, 224)
c = sum_a[:, :, None] * b  # 广播后形状为(224, 224, f)

方法2：向量化逐通道计算（通用场景） #

如果后续需要修改操作逻辑（比如替换求和为其他通道运算），可以用广播实现逐通道元素乘法后再聚合：

# 将a扩展为(224, 224, 1, 3)，b扩展为(224, 224, f, 1)，广播相乘后求和最后一维
c = tf.reduce_sum(a[:, :, None, :] * b[:, :, :, None], axis=-1)

方法3：使用Einstein求和（可读性强） #

用tf.einsum可以直观表达张量维度间的运算关系，对应需求的表达式为hwd,hwf->hwf（h=高度，w=宽度，d=a的通道数，f=b的最后一维数）：

c = tf.einsum('hwd,hwf->hwf', a, b)

关于你提到的Conv2D等效性 #

普通的Conv2D(filters=f, kernel_size=1)是所有空间位置共享同一组3×f的权重，而你的操作是每个空间位置使用独立的f个标量权重（即b的对应位置值），因此两者并不完全等价。若将你的操作视为“逐位置的1×1卷积”，虽可通过动态权重实现，但效率远低于前面的广播方法，不推荐使用。

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