Multiply outputs from two Conv2D layers in TensorFlow 2+

Question

I'm trying to implement SOLO architecture for instance segmentation in TensorFlow (Decoupled version).

https://arxiv.org/pdf/1912.04488.pdf

Right now, I need to compute the loss function and multiply each output map from first conv2d layer with output maps of second layer.

xi = Conv2D(…)(input) # output is (batch, None, None, 24)
yi = Conv2D(…)(input) # output is (batch, None, None, 24)

I need to multiply each output maps (element wise) xi with yi in a way to get output with (batch, None, None, 24*24) . I need element-wise multiplication of one output feature map (from first conv2d) with all from second conv2d layer and so on. Thats why 24 * 24 .

I try to do this with for cycles but get error:

OperatorNotAllowedInGraphError: iterating over tf.Tensor is not allowed: 
                                                    AutoGraph did convert this function.

Any advice to achieve this with some TF2 operation?

Answer 1

It is basically the outer product of the last dimension followed by collapsing the last 2 dimensions. A short way to express this operation is to use tf.repeat and tf.tile .

#channel_dims is the length of the last dimension, i.e. 24 in your question 
@tf.function
def outerprodflatten(x,y,channel_dims):
  return tf.repeat(x,channel_dims,-1)*tf.tile(y,[1,1,1,channel_dims])

To use this in functional API, you would also have to define a custom keras layer or lambda layer, ie

class Outerprodflatten_layer(tf.keras.layers.Layer):
    def __init__(self, channel_dims):
        self.channel_dims = channel_dims
        super().__init__()

    def call(self, inputs):
        return tf.repeat(inputs[0],self.channel_dims,-1)*tf.tile(inputs[1],[1,1,1,self.channel_dims])

x = Conv2D(…)(input) # output is (batch, None, None, 24)
y = Conv2D(…)(input) # output is (batch, None, None, 24)

out=Outerprodflatten_layer(24)([x,y])

Answer 2

You can use the Multiply layer to do this

multiplied = tensorflow.keras.layers.Multiply()([xi,yi])

This layer multiplies (element-wise) a list of inputs.
It takes as input a list of tensors, all of the same shape, and returns a single tensor (also of the same shape).