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 ）

依此類推，直到第f次乘法。 結果將是形狀為 224x224 的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()

Output：

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>