any way to rescale and stack features maps with different shape in tensorflow?

user_10708392 0
Viewed 152

I intend to use the concepts of skip connection in my experiment. Basically, in my pipeline, features maps that comes after Conv2D are going to be stacked or concatenated. But, features maps are in different shape and try to stack them together into one tensor gave me error. Does anyone knows any possible way of doing this correctly in tensorflow? Any thoughts or ideas to make this happen? Thanks

idea flowchart

here is the pipeline flowchart I want to do it:

enter image description here

my case is little different because I got extra building block is used after Conv2D and its output now is feature maps of 15x15x64 and so on. I want to stack those features map into one then use it to Conv2D again.

my attempt:

this is my reproducible attempt:

import tensorflow as tf
from tensorflow.keras.layers import Dense, Dropout, Activation, Conv2D, Flatten, MaxPool2D, BatchNormalization


inputs = tf.keras.Input(shape=(32, 32, 3))
x = inputs
x = Conv2D(32, (3, 3), input_shape=(32,32,3))(x)
x = BatchNormalization(axis=-1)(x)
x = Activation('relu')(x)
fm1 = MaxPooling2D(pool_size=(2,2))(x)

x = Conv2D(32,(3, 3), input_shape=(15,15,32))(fm1)
x = BatchNormalization(axis=-1)(x)
x = Activation('relu')(x)
fm2 = MaxPooling2D(pool_size=(2,2))(x)

concatted = tf.keras.layers.Concatenate(axis=1)([fm1, fm2])

but this way I ended up with following error: ValueError: A Concatenate layer requires inputs with matching shapes except for the concat axis. Got inputs shapes: [(None, 15, 15, 32), (None, 6, 6, 32)]. I am not sure what would be correct way to stack features maps with different shape. How can we make this right? Any possible thoughts?

desired output

in my actual model, I got shape of features maps are TensorShape([None, 15, 15, 128]) and TensorShape([None, 6, 6, 128]). I need to find way to merge them or stack them into one. Ideally, shape of concatenated or stacked feature maps' shape would be: [None, 21,21,128]. Is there any way of stacking them into one? Any idea?

tensorflow python
1 Answers

What you're trying to achieve doesn't work mathematically. Let me illustrate. Take the simple 1D problem (like 1D convolution). You have a (None, 64, 128) (fm1) sized output and a (None, 32, 128) (fm2) output that you want to concatenate. Then,

concatted = tf.keras.layers.Concatenate(axis=1)([fm1, fm2])

works totally fine, giving you an output of size (None, 96, 128).

Let's come to the 2D problem. Now you got two tensors (None, 15, 15, 128) and (None, 6, 6, 128) and want to end up with a (None, 21, 21, 128) sized output. Well the math doesn't work here. To understand why, reduce this to 1D format. Then you got

fm1 -> (None, 225, 128)
fm2 -> (None, 36, 128)

By concat you get,

concatted -> (None, 261, 128)

If the math works you should get (None, 441, 128) which is reshape-able to (None, 21, 21, 128). So this cannot be achieved unless you pad the edges of the smaller with 441-261 = 180 on the reshaped tensor. And then reshape it to the desired shape. Following is an example of how you can do it,

concatted = tf.keras.layers.Lambda(
    lambda x: K.reshape(
                  K.concatenate(
                      [K.reshape(x[0], (-1, 225, 128)), 
                       tf.pad(
                           K.reshape(x[1], (-1, 36, 128)), [(0,0), (0, 180), (0,0)]
                       )
                      ], axis=1
                  ), (-1, 21, 21, 128))
              )([fm1, fm2])

Important: But I can't guaranttee the performance of your model this just solves your problem mathematically. In a machine learning perspective, I wouldn't advice this. Best way would be making sure the outputs are compatible in sizes for concatenation. Few ways would be,

  • Not reduce the size of convolution outputs (stride = 0 and padding='same')
  • Use transpose convolution operation to size-up the smaller one
user_16990750
Related