keras, TensorFlow中加入注意力机制

快来打我* 2023-06-25 03:12 19阅读 0赞

### 第一步：找到要修改文件的源代码 ###

## 在里面添加通道注意力机制和空间注意力机制 ##

##### 所需库 #####

from keras.layers import GlobalAveragePooling2D, GlobalMaxPooling2D, Reshape, Dense, multiply, Permute, Concatenate, Conv2D, Add, Activation, Lambda
    from keras import backend as K
    from keras.activations import sigmoid

#### 通道注意力机制 ####

def channel_attention(input_feature, ratio=8):
    	
    	channel_axis = 1 if K.image_data_format() == "channels_first" else -1
    	channel = input_feature._keras_shape[channel_axis]
    	
    	shared_layer_one = Dense(channel//ratio,
    							 kernel_initializer='he_normal',
    							 activation = 'relu',
    							 use_bias=True,
    							 bias_initializer='zeros')
    
    	shared_layer_two = Dense(channel,
    							 kernel_initializer='he_normal',
    							 use_bias=True,
    							 bias_initializer='zeros')
    	
    	avg_pool = GlobalAveragePooling2D()(input_feature)    
    	avg_pool = Reshape((1,1,channel))(avg_pool)
    	assert avg_pool._keras_shape[1:] == (1,1,channel)
    	avg_pool = shared_layer_one(avg_pool)
    	assert avg_pool._keras_shape[1:] == (1,1,channel//ratio)
    	avg_pool = shared_layer_two(avg_pool)
    	assert avg_pool._keras_shape[1:] == (1,1,channel)
    	
    	max_pool = GlobalMaxPooling2D()(input_feature)
    	max_pool = Reshape((1,1,channel))(max_pool)
    	assert max_pool._keras_shape[1:] == (1,1,channel)
    	max_pool = shared_layer_one(max_pool)
    	assert max_pool._keras_shape[1:] == (1,1,channel//ratio)
    	max_pool = shared_layer_two(max_pool)
    	assert max_pool._keras_shape[1:] == (1,1,channel)
    	
    	cbam_feature = Add()([avg_pool,max_pool])
    	cbam_feature = Activation('hard_sigmoid')(cbam_feature)
    	
    	if K.image_data_format() == "channels_first":
    		cbam_feature = Permute((3, 1, 2))(cbam_feature)
    	
    	return multiply([input_feature, cbam_feature])

#### 空间注意力机制 ####

def spatial_attention(input_feature):
    	kernel_size = 7
    	
    	if K.image_data_format() == "channels_first":
    		channel = input_feature._keras_shape[1]
    		cbam_feature = Permute((2,3,1))(input_feature)
    	else:
    		channel = input_feature._keras_shape[-1]
    		cbam_feature = input_feature
    	
    	avg_pool = Lambda(lambda x: K.mean(x, axis=3, keepdims=True))(cbam_feature)
    	assert avg_pool._keras_shape[-1] == 1
    	max_pool = Lambda(lambda x: K.max(x, axis=3, keepdims=True))(cbam_feature)
    	assert max_pool._keras_shape[-1] == 1
    	concat = Concatenate(axis=3)([avg_pool, max_pool])
    	assert concat._keras_shape[-1] == 2
    	cbam_feature = Conv2D(filters = 1,
    					kernel_size=kernel_size,
    					activation = 'hard_sigmoid',
    					strides=1,
    					padding='same',
    					kernel_initializer='he_normal',
    					use_bias=False)(concat)
    	assert cbam_feature._keras_shape[-1] == 1
    	
    	if K.image_data_format() == "channels_first":
    		cbam_feature = Permute((3, 1, 2))(cbam_feature)
    		
    	return multiply([input_feature, cbam_feature])

#### 构建CBAM ####

def cbam_block(cbam_feature,ratio=8):
    	"""Contains the implementation of Convolutional Block Attention Module(CBAM) block.
    	As described in https://arxiv.org/abs/1807.06521.
    	"""
    	
    	cbam_feature = channel_attention(cbam_feature, ratio)
    	cbam_feature = spatial_attention(cbam_feature, )
    	return cbam_feature

##### 在相应的位置添加CBAM #####

inputs = x
      residual = layers.Conv2D(filter, kernel_size = (1, 1), strides = strides, padding = 'same')(inputs)
      residual = layers.BatchNormalization(axis = bn_axis)(residual)
      cbam = cbam_block(residual)
      x = layers.add([x, residual, cbam])

这样就在任意位置加入了注意力机制啦。