team7/talk_with_wind_test_something

1import torch2import torch.nn as nn3import torch.nn.functional as F4from torch import Tensor5 6from efficientat.models.utils import collapse_dim7 8 9class MultiHeadAttentionPooling(nn.Module):10    """Multi-Head Attention as used in PSLA paper (https://arxiv.org/pdf/2102.01243.pdf)11    """12    def __init__(self, in_dim, out_dim, att_activation: str = 'sigmoid',13                 clf_activation: str = 'ident', num_heads: int = 4, epsilon: float = 1e-7):14        super(MultiHeadAttentionPooling, self).__init__()15 16        self.in_dim = in_dim17        self.out_dim = out_dim18        self.num_heads = num_heads19        self.epsilon = epsilon20 21        self.att_activation = att_activation22        self.clf_activation = clf_activation23 24        # out size: out dim x 2 (att and clf paths) x num_heads25        self.subspace_proj = nn.Linear(self.in_dim, self.out_dim * 2 * self.num_heads)26        self.head_weight = nn.Parameter(torch.tensor([1.0 / self.num_heads] * self.num_heads).view(1, -1, 1))27 28    def activate(self, x, activation):29        if activation == 'linear':30            return x31        elif activation == 'relu':32            return F.relu(x)33        elif activation == 'sigmoid':34            return torch.sigmoid(x)35        elif activation == 'softmax':36            return F.softmax(x, dim=1)37        elif activation == 'ident':38            return x39 40    def forward(self, x) -> Tensor:41        """x: Tensor of size (batch_size, channels, frequency bands, sequence length)42        """43        x = collapse_dim(x, dim=2)  # results in tensor of size (batch_size, channels, sequence_length)44        x = x.transpose(1, 2)  # results in tensor of size (batch_size, sequence_length, channels)45        b, n, c = x.shape46 47        x = self.subspace_proj(x).reshape(b, n, 2, self.num_heads, self.out_dim).permute(2, 0, 3, 1, 4)48        att, val = x[0], x[1]49        val = self.activate(val, self.clf_activation)50        att = self.activate(att, self.att_activation)51        att = torch.clamp(att, self.epsilon, 1. - self.epsilon)52        att = att / torch.sum(att, dim=2, keepdim=True)53 54        out = torch.sum(att * val, dim=2) * self.head_weight55        out = torch.sum(out, dim=1)56        return out57