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

Question

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

我有一个形状为224x224x3的图像张量a和形状为224x224xf的张量b。我想将(元素方面)a乘以每个由f切片的b的二维矩阵，得到形状为224x224xf的矩阵c。

因此，例如，第一次乘法将按如下方式进行:

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

(广播+求和，结果为形状224x224)

等等，直到第f次乘法。结果是形状为224x224的f矩阵在形状为224x224xf的c矩阵中聚集。

我将非常感谢如何使用tensorflow功能的帮助。

编辑:我意识到我想做的是相当于一个Conv2D操作与kernel_size=1和过滤器=f。也许能帮上忙。

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

Answer 2

你可以分别在倒数第二个维度和最后一个维度展开这两个张量，然后利用广播。

tf.reduce_sum(tf.expand_dims(a, axis=-2) * tf.expand_dims(b[..., :f+1], axis=-1), axis=-1)

证明这产生正确的结果

a = tf.random.uniform(shape=(224,224,3))
b = tf.random.uniform(shape=(224,224,10))
f = 4
ref = None
for i in range(f+1):
if ref is None:
ref = tf.reduce_sum(a * b[...,i][...,None], axis=-1)[...,None]
else:
ref = tf.concat([ref, tf.reduce_sum(a * b[...,i][...,None], axis=-1)[...,None]], axis=-1)
tf.reduce_all(tf.reduce_sum(tf.expand_dims(a, axis=-2) * tf.expand_dims(b[..., :f+1], axis=-1), axis=-1) == ref)
<tf.Tensor: shape=(), dtype=bool, numpy=True>