3D matrix multiplication in tensorflow: AAB and AAB matrices to get a new AAB matrix

Viewed 72

Suppose I have two images with dimensions of 32x32x3 (number of channels=3). I want to multiply them (like "matmul" function) on the first and the second dimensions for each of these 3 channels in Tensorflow to get a new 32x32x3 image. Can someone help me with this? Something like this loop:

#x.shape=(32,32,3)
#y.shape=(32,32,3)
a = np.zeros((x.shape[-3], x.shape[-2], x.shape[-1],), dtype='float32')
for i in range(a.shape[-1]):
    a[:, :, i] = tf.matmul(x[:, :, i], y[:, :, i])
a = tf.convert_to_tensor(a, dtype=tf.float32)

but I was wondering there is a more efficient way to do this?

1 Answers

Actually, I found the answer. The matmul works also for 3d arrays. However, the features (channels) need to be first in matmul function. So we need to use tf.transpose if channels are placed in the last dimension as follow:

x=tf.transpose(x, perm=[2, 0, 1])
y=tf.transpose(y, perm=[2, 0, 1])

a=tf.matmul(x,y)
a=tf.transpose(a, perm=[1, 2, 0])

It gives the same result as the loop I wrote above.

Related