El código pytorch implementa SimAM del mecanismo de atención.

Mecanismo de atención SimAM

SimAM (Mecanismo de atención basado en la similitud) es un mecanismo de atención para calcular la similitud en los datos de secuencia.
Dirección del artículo: http://proceedings.mlr.press/v139/yang21o/yang21o.pdf
En SimAM, dada una secuencia de consulta Q y una secuencia de par clave-valor K, el mecanismo de atención calcula la relación entre la secuencia de consulta y la clave similitud de secuencia para determinar el peso de la atención. Específicamente, SimAM usa la similitud del coseno para medir la similitud entre las secuencias de consulta y las secuencias clave.
El proceso de cálculo de SimAM es el siguiente:
(1) Para cada elemento q en la secuencia de consulta Q, calcule su similitud de coseno con cada elemento k en la secuencia clave K.
(2) Para cada elemento de secuencia de consulta q, el peso de atención se obtiene normalizando su similitud con el elemento de secuencia clave k.
(3) La secuencia de valores V se pondera y suma usando pesos de atención para obtener la representación de atención final.

el código se muestra a continuación:

import torch
import torch.nn as nn


class SimAM(torch.nn.Module):
    def __init__(self, e_lambda=1e-4):
        super(SimAM, self).__init__()

        self.activaton = nn.Sigmoid()
        self.e_lambda = e_lambda

    def __repr__(self):
        s = self.__class__.__name__ + '('
        s += ('lambda=%f)' % self.e_lambda)
        return s

    @staticmethod
    def get_module_name():
        return "simam"

    def forward(self, x):
        b, c, h, w = x.size()

        n = w * h - 1

        x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)
        y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5

        return x * self.activaton(y)


if __name__ == '__main__':
    input = torch.randn(3, 64, 7, 7)
    model = SimAM()
    outputs = model(input)
    print(outputs.shape)