Tensorflow 中矩阵的逐元素乘法：如何避免 for 循环

Question

我想在 tensorflow (TF 2.10) 中进行以下乘法运算，但我不确定如何进行。

我有一个图像张量a，形状为224x224x3，还有一个张量b，形状为224x224xf。我想将（按元素）a乘以b的每个二维矩阵，并用f切片得到形状为224x224xf的矩阵c。

因此，例如，第一次乘法将按如下方式完成：

tf.reduce_sum(a * b[:,:,0][:,:,None],axis=-1)

（广播+求和，结果是形状224x224）

依此类推，直到 fth 乘法。结果将是f形状为224x224的矩阵在c形状为224x224xf的矩阵中的聚合。

我将非常感谢有关如何使用 tensorflow 功能执行此操作的帮助。

编辑：我意识到我想做的相当于 kernel_size=1 和 filters=f 的 Conv2D 操作。也许它可以帮助。

Answer 1

您可以将 a 的每个通道与 b 相乘，然后求和：

X = a[:,:,0:1] * b + a[:,:,1:2] * b + a[:,:,2:3] * b

X 的形状是(224, 224, f)，它会给出与乘法相同的结果：

(X[:, :, 0] == tf.reduce_sum(a * b[:, :, 0][:, :, None], axis=-1)).numpy().all()

输出：

True

以下给出了略有不同的结果，我猜是因为浮点舍入：

tf.reduce_sum(a, axis=-1, keepdims=True) * b