How can I change the shape of a feature map with dimension 1x1xC into NxNxC dimension feature map in CNN?

Viewed 194
image_size = 256 #image size which is input shape of the densenet121
base_model=get_base_model('densenet121',image_size) #calling densenet121 using transfer learning
base_in=base_model.input #base_in.shape = (None, 256, 256, 3)
base_out=base_model.output #base_out.shape = (None, 8, 8, 1024)

#function which performs attention which is to be added on the densenet121
def BAM_Channel_Attention(inputs):
  shape=K.int_shape(inputs)
  x=tf.keras.layers.AveragePooling2D(pool_size=(shape[1],shape[2]))(inputs)
  dense = layers.Dense(1280, activation="relu")(x)
  reduction = Conv2D(shape[3]/16,1, padding='same')(dense)
  dense = layers.Dense(1280, activation="relu")(reduction)
  out = tf.keras.layers.BatchNormalization()(dense)
  return out

the above function is called here:

#calling the attention function with the output of the densenet121
x=BAM_Channel_Attention(base_out) #

print(x.shape)

I get the output (None, 1, 1, 1280). How can I change the feature map from 1x1x1280 to say 8x8x1280?

1 Answers

The simplest, but perhaps not the most practical option would be to use tf.repeat on the BAM_Channel_Attention output:

import tensorflow as tf

batch_size = 2
data = tf.random.normal((batch_size, 1, 1, 1280))
output = tf.repeat(tf.repeat(data, repeats=8, axis=1), repeats=8, axis=2)
print(output.shape)
(2, 8, 8, 1280)

A probably better approach is to change the pool_size in the AveragePooling2D layer in the BAM_Channel_Attention function, since it is downsampling your spatial dimensions from (8, 8) to (1,1). You could change your pool_size from (8, 8) to (1, 1):

import tensorflow as tf

batch_size = 2
inputs = tf.keras.layers.Input((8, 8, 1024))
x = tf.keras.layers.AveragePooling2D(pool_size=(1,1))(inputs)
dense = tf.keras.layers.Dense(1280, activation="relu")(inputs)
reduction = tf.keras.layers.Conv2D(1024/16, 1, padding='same')(dense)
dense = tf.keras.layers.Dense(1280, activation="relu")(reduction)
outputs = tf.keras.layers.BatchNormalization()(dense)

model = tf.keras.Model(inputs, outputs)
print(model(tf.random.normal((batch_size, 8, 8, 1024))).shape)
(2, 8, 8, 1280)

Or completely omit the AveragePooling2D layer and you will also have the result you want.

Related