Element-wise multiplication of matrices in Tensorflow : how to avoid for loop-CodePudding

I want to do the following multiplication in tensorflow (TF 2.10), but I'm not sure how to.

I have an image tensor a, which is of shape 224x224x3 and a tensor b, which is of shape 224x224xf. I want to multiply (element-wise) a by each 2D matrix of b sliced by f to get a matrix c of shape 224x224xf.

So for example, the 1st multiplication would be done as follows:

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

(broadcasting summation, result is shape 224x224)

and so on until the fth multiplication. Result would be the aggregation of f matrices of shape 224x224 in c matrix of shape 224x224xf.

I would greatly appreciate help on how to do this using tensorflow functionality.

EDIT: I realize that what I want to do is equivalent to a Conv2D operation with kernel_size=1 and filters=f. Maybe it can help.

CodePudding user response：

You could multiply each channel of a with b and then sum:

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

The shape of X is (224, 224, f) and it will give the same results as your multiplications:

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

Output:

True

The following gives slightly different results, I guess because of floating point rounding:

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

CodePudding user response：

You can expand the two tensors in next-to-last and last dimension, respectively, then take advantage of broadcasting.

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

Proof that this produces correct result

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>