Implementación del mecanismo de atención CBAM

1. Introducción a CBAM

Documento: CBAM:
Módulo de atención de bloques convolucionales
propone un módulo de atención simple pero efectivo CBAM, dado un mapa de características intermedias, infiere el peso de la atención a lo largo de las dos dimensiones de espacio y canal, y luego lo multiplica con el mapa de características original Realice ajustes adaptativos a las funciones. Dado que CBAM es un módulo ligero de uso general, se puede integrar sin problemas en cualquier arquitectura de CNN con una sobrecarga adicional insignificante, y se puede entrenar de un extremo a otro junto con la CNN básica. Después de integrar CBAM en diferentes modelos en diferentes conjuntos de datos de clasificación y detección, el rendimiento del modelo se ha mejorado constantemente, lo que demuestra su amplia aplicabilidad.
CBAM se puede dividir en dos partes: módulo de atención del canal y módulo de atención espacial, como se muestra en la figura siguiente.

• Módulo de atención del canal: preste atención a las características que son significativas.
  Ingrese un mapa de características F como H W C (de hecho, puede haber un lote o NHWC). Primero, realice una agrupación máxima de espacio global y una agrupación promedio, y obtenga dos Un descriptor 1 1 C. Luego, envíelos a MLP (que contiene una capa oculta), el número de neuronas en la primera capa es C / r, y el número de neuronas en la segunda capa es C. Esta red neuronal es compartida (no sé si se comparte entre dos descriptores o dos capas de redes neuronales, que deberían ser las últimas). Luego aplique la suma de elementos a los dos vectores de salida y obtenga el coeficiente de peso Mc a través de la función sigmoidea. Finalmente, multiplique Mc con la característica original F para obtener una nueva característica F '
• Módulo de atención espacial: Preste atención a dónde las características son significativas.
  Ingrese el mapa de características F'as H W C (también puede haber un lote). Primero, se obtienen dos descriptores de H W 1 a través de la combinación promedio de canales y la combinación máxima . Y empalmado por canal. Luego, a través de una capa convolucional 7 * 7 y una función de activación sigmoidea, se obtiene el coeficiente de peso Ms. Finalmente, multiplique Ms y F 'para obtener la característica de atención final.

2. Implementación del código

Después de leer el proceso anterior, es fácil comprender la implementación del código. Aquí hay dos piezas de código. Uno es hacer referencia al código del blogger y el otro es Github: kobiso / CBAM-tensorflow .
Codigo uno:

def combined_static_and_dynamic_shape(tensor):
    """Returns a list containing static and dynamic values for the dimensions.  Returns a list of static 
    and dynamic values for shape dimensions. This is  useful to preserve static shapes when available in reshape operation.  
    Args:    tensor: A tensor of any type.  
    Returns:    A list of size tensor.shape.ndims containing integers or a scalar tensor.  """
    static_tensor_shape = tensor.shape.as_list()
    dynamic_tensor_shape = tf.shape(tensor)
    combined_shape = []
    for index, dim in enumerate(static_tensor_shape):
        if dim is not None:
            combined_shape.append(dim)
        else:
            combined_shape.append(dynamic_tensor_shape[index])
    return combined_shape


def convolutional_block_attention_module(feature_map, index, reduction_ratio = 0.5):
    """CBAM：convolutional block attention module
    Args:
        feature_map:input feature map
        index:the index of the module
        reduction_ratio:output units number of first MLP layer:reduction_ratio * feature map
    Return:
        feature map with channel and spatial attention"""

    with tf.variable_scope("cbam_%s" % (index)):
        feature_map_shape = combined_static_and_dynamic_shape(feature_map)
        # channel attention module
        channel_avg_weights = tf.nn.avg_pool(value=feature_map,
                                             ksize=[1, feature_map_shape[1], feature_map_shape[2], 1],
                                             strides=[1, 1, 1, 1],
                                             padding='VALID')  # global average pool
        channel_max_weights = tf.nn.max_pool(value=feature_map,
                                             ksize=[1, feature_map_shape[1], feature_map_shape[2], 1],
                                             strides=[1, 1, 1, 1],
                                             padding='VALID')
        channel_avg_reshape = tf.reshape(channel_avg_weights,
                                         [feature_map_shape[0], 1, feature_map_shape[3]])
        channel_max_reshape = tf.reshape(channel_max_weights,
                                         [feature_map_shape[0], 1, feature_map_shape[3]])
        channel_w_reshape = tf.concat([channel_avg_reshape, channel_max_reshape], axis=1)

        fc_1 = tf.layers.dense(inputs=channel_w_reshape,
                               units=feature_map_shape[3] * reduction_ratio,
                               name="fc_1",
                               activation=tf.nn.relu)
        fc_2 = tf.layers.dense(inputs=fc_1,
                               units=feature_map_shape[3],
                               name="fc_2",
                               activation=None)
        channel_attention = tf.reduce_sum(fc_2, axis=1, name="channel_attention_sum")
        channel_attention = tf.nn.sigmoid(channel_attention)
        channel_attention = tf.reshape(channel_attention,
                                       shape=[feature_map_shape[0], 1, 1, feature_map_shape[3]])
        feature_map_with_channel_attention = tf.multiply(feature_map, channel_attention)
        # saptial attention module
        # 通道平均池化，格式NWHC
        channel_wise_avg_pooling = tf.reduce_mean(feature_map_with_channel_attention, axis=3)
        channel_wise_avg_pooling = tf.reshape(channel_wise_avg_pooling,
                                              shape=[feature_map_shape[0], feature_map_shape[1],
                                                     feature_map_shape[2], 1]) # shape=[batch, H, W, 1]
        # 通道最大池化
        channel_wise_max_pooling = tf.reduce_max(feature_map_with_channel_attention, axis=3)
        channel_wise_max_pooling = tf.reshape(channel_wise_max_pooling,
                                              shape=[feature_map_shape[0], feature_map_shape[1],
                                                     feature_map_shape[2], 1])
        # 按通道拼接
        channel_wise_pooling = tf.concat([channel_wise_avg_pooling, channel_wise_max_pooling], axis=3)
        spatial_attention = slim.conv2d(channel_wise_pooling, 1, [7, 7],
                                        padding='SAME',
                                        activation_fn=tf.nn.sigmoid,
                                        scope="spatial_attention_conv")
        feature_map_with_attention = tf.multiply(feature_map_with_channel_attention, spatial_attention)
        return feature_map_with_attention

Cuando este código se implementa en el módulo de atención del canal, los dos descriptores 1 1 C se empalman primero y luego se ingresan en el MLP, y los pesos de las dos capas MLP no se comparten, lo que se siente un poco problemático. Debido al límite de palabras, se colocará otro código en el próximo blog.
Puse este módulo en una de mis redes neuronales de cuatro capas para la clasificación MNIST manuscrita, pero no creo que la precisión haya mejorado mucho e intuitivamente siento que funciona más lentamente. En el documento, al insertar CBAM en algunas redes grandes, siento que la mejora del rendimiento no es particularmente grande. . .