Can we use squeeze-and-excite attention block on pretrained network? (only train SENET block, will it work?)

Viewed 14

I am using pre-trained VGG19 and applying squeeze-and-excitation attention (SENET) module as mentioned in the Squeeze and Excitaion Network, in between pre-trained VGG19 layer, and only training the last-fully connected layer and SENET block.

Will it work? Am I doing it right, or should I train the complete network? (My analogy behind training only the SENET is that it will suppress or enhance particular feature channels by learning attention weights during the training.)

I've attached my code below for reference:

import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import Activation, Conv2D, Input, BatchNormalization, Reshape, GlobalAveragePooling2D, AveragePooling2D, GlobalMaxPooling2D, Flatten , ReLU, Layer, Dense
from tensorflow.keras.activations import sigmoid, softmax, relu, tanh
from tensorflow.keras import Sequential

class SENET_Attn(Layer):
    """
    Channel Attention Block as reported in SENET
    """
    def __init__(self,out_dim, ratio, layer_name="SENET"):
        super(SENET_Attn, self).__init__()
        self.out_dim = out_dim
        self.ratio = ratio
        self.layer_name = layer_name
            
    def build(self, ratio, layer_name="SENET"):
        self.Global_Average_Pooling = GlobalAveragePooling2D(keepdims= True)
        self.Fully_connected_1_1 = Dense(units= self.out_dim/self.ratio, name=self.layer_name+'_fully_connected1',
                                kernel_initializer="glorot_uniform")
        self.Relu = ReLU()
        self.Fully_connected_2 = Dense(units=self.out_dim, name=layer_name+'_fully_connected2', activation = "tanh")
        self.Sigmoid = Activation("sigmoid")
        
    def call(self, inputs):
        inputs = tf.cast(inputs, dtype = "float32")
        squeeze = self.Global_Average_Pooling(inputs)
        excitation = self.Fully_connected_1_1(squeeze)
        excitation = self.Relu(excitation)
        excitation = self.Fully_connected_2(excitation)
        excitation =  self.Sigmoid(excitation)
        excitation = tf.reshape(excitation, [-1,1,1,self.out_dim])
        
        scale = inputs * excitation
        return scale
Vgg = VGG19(include_top = False,input_shape = (50,50,3))

# SENET-Attn VGG19
ratio =16

input_layer = Input(shape=(50,50,3))
out = Vgg.layers[1](input_layer) # Block one of VGG19
out = Vgg.layers[2](out)
out = Vgg.layers[3](out)

out = Vgg.layers[4](out) # Block two of VGG19
out = Vgg.layers[5](out)
out = Vgg.layers[6](out)

#out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[7](out) # Block three of VGG19
out = Vgg.layers[8](out)
out = Vgg.layers[9](out)
out = Vgg.layers[10](out)
out = Vgg.layers[11](out)
 
out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[12](out) # Block four of VGG19
out = Vgg.layers[13](out)
out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[14](out)
out = Vgg.layers[15](out)
out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[16](out)

out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[17](out) #Block five of VGG19

out = Vgg.layers[18](out)
out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[19](out)
out = Vgg.layers[20](out)
out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
out = Vgg.layers[21](out)

out = SENET_Attn(out.shape[-1], ratio, )(out) # SENET Attention
flatten = Flatten()(out)
out = Dense(100)(flatten)
out = tf.keras.layers.ReLU()(out)
out = Dense(100)(out)
out = tf.keras.layers.ReLU()(out)
out = Dense(2)(out)
out = tf.keras.layers.Softmax()(out)

model = Model(input_layer, out)
model.compile(optimizer= tf.keras.optimizers.SGD(), loss= tf.keras.losses.CategoricalCrossentropy()
              , metrics=['accuracy'])
model.summary()
0 Answers
Related