CODEWITHSUNDEEP
tensorflowfor-loopmatrixkerastensor
Jane Sully
Jane Sully

Reputation: 3327

Replacing for loop with tensor multiplication in tensorflow

I want to do the following normalization in tensorflow, but I'm not sure how to.

I have a tensor a, which is of shape (c, c) and let's say c=16 in this case. I have another tensor of shape b, which is of shape (w, w, c) and w = 32, c=16 (so b's dimensions are height w, width w, and depth/channels c). I'm ignoring the batch dimensions in both cases, but you can assume the a and b have the same batch dimensions/number of samples.

I would like to do something like the following, ideally using tensor notation/operations and not using this for loop:

for i in range(c): # loop over the channels
    y[:,:,i] = b[:,:,i]/(sum_j[a[j,i] * b[:,:,j]])

In other words, we're normalizing b by a weighted sum over j a[j,i]*b[:,:,j], where j is indexing the channel. We do this for each channel i to produce y[:,:,i]. I would greatly appreciate help on how to do this using tensorflow functionality. Thanks!

Upvotes: 1

Views: 510

Answers (1)

o-90
o-90

Reputation: 17585

Let c = 2 and w = 3 for illustration.

import tensorflow as tf


a = tf.constant([
    [1, 2],
    [3, 4]], dtype=tf.float32)

b = tf.constant([
    [[11, 15],
     [4, 3],
     [3, 7]],
    [[22, 5],
     [15, 20],
     [1, 4]],
    [[16, 14],
     [3, 22],
     [2, 8]]], tf.float32)

a = a[:, None, None]
print(a.shape)
# (2, 1, 1, 2)

b_extended = tf.tile(b[None], [2, 1, 1, 1])
print(b_extended.shape)
# (2, 3, 3, 2)

norm = tf.math.reduce_sum(a * b_extended, -1)
norm = tf.transpose(norm, (1, 2, 0))

out = b / norm

print(out)
# tf.Tensor(
# [[[0.2682927  0.16129032]
#   [0.4        0.125     ]
#   [0.1764706  0.1891892 ]]
# 
#  [[0.6875     0.05813954]
#   [0.27272728 0.16      ]
#   [0.11111111 0.21052632]]
# 
#  [[0.36363637 0.13461539]
#   [0.06382979 0.22680412]
#   [0.11111111 0.21052632]]], shape=(3, 3, 2), dtype=float32)

Let's test with numpy and see if this is indeed what we want.

import numpy as np


a = a.numpy()
b = b.numpy()

y = np.zeros((3, 3, 2))

for i in range(2):
    y[:, :, i] = b[:, :, i] / (a[i] * b).sum(-1)

np.testing.assert_array_almost_equal(y, out.numpy())
# passes

Edit

You can just wrap it in a tensorflow.keras.layers.Layer and make your own layer and you'll need to account for batches now when you are expanding the inputs.

import tensorflow as tf
from tensorflow.keras import layers


class CustomNormLayer(layers.Layer):
    """``CustomNormLayer``."""
    def __init__(self):
        super().__init__()
        
    def call(self, x, y):
        # extract `c`
        c = tf.shape(x)[1]  
        # expand x
        x = x[:, :, None, None]
        # ----^ (batch)
        # tile y
        y_extended = tf.tile(y[:, None], (1, c, 1, 1, 1))
        # ---------------------^ (batch)
        # sum norm
        norm = tf.math.reduce_sum(x * y_extended, -1)
        # permute
        norm = tf.transpose(norm, (0, 2, 3, 1))
        return y / norm

Let's make a model and try it out

a = tf.constant([
    [1, 2],
    [3, 4]], dtype=tf.float32)

b = tf.constant([
    [[11, 15],
     [4, 3],
     [3, 7]],
    [[22, 5],
     [15, 20],
     [1, 4]],
    [[16, 14],
     [3, 22],
     [2, 8]]], tf.float32)

# put data into a `tf.data` object
batch_size = 5
dataset = tf.data.Dataset.from_tensors((a, b))
dataset = dataset.repeat()
dataset = dataset.batch(batch_size)
train_iter = iter(dataset)

# define layers
xa = tf.keras.Input((2, 2))
xb = tf.keras.Input((3, 3, 2))
out = CustomNormLayer()(xa, xb)

# define model
model = tf.keras.Model([xa, xb], out)

# put some data in
data = next(train_iter)
model(data)[0]
# <tf.Tensor: shape=(3, 3, 2), dtype=float32, numpy=
# array([[[0.2682927 , 0.16129032],
#         [0.4       , 0.125     ],
#         [0.1764706 , 0.1891892 ]],
# ...

Upvotes: 1

Related Questions