le code pytorch implémente le mécanisme d'attention SimAM

Mécanisme d'attention SimAM

SimAM (Similarity-based Attention Mechanism) est un mécanisme d'attention pour calculer la similarité dans les données de séquence.
Adresse papier : http://proceedings.mlr.press/v139/yang21o/yang21o.pdf
Dans SimAM, étant donné une séquence de requête Q et une séquence de paire clé-valeur K, le mécanisme d'attention calcule la relation entre la séquence de requête et la clé similitude de séquence pour déterminer le poids de l'attention. Plus précisément, SimAM utilise la similarité cosinus pour mesurer la similarité entre les séquences de requêtes et les séquences de clés.
Le processus de calcul de SimAM est le suivant :
(1) Pour chaque élément q dans la séquence de requête Q, calculez sa similarité cosinus avec chaque élément k dans la séquence de clé K.
(2) Pour chaque élément de séquence de requête q, le poids d'attention est obtenu en normalisant sa similarité avec l'élément de séquence de clé k.
(3) La séquence de valeurs V est pondérée et additionnée à l'aide de poids d'attention pour obtenir la représentation finale de l'attention.

code afficher comme ci-dessous:

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)