SelfAttention

SelfAttention.py

import torch
import torch.nn as nn
  
class SelfAttention(nn.Module):
  
    def __init__(self, input_dim, output_dim, dropout=0.1):
        super(SelfAttention, self).__init__()

        self.q = nn.Linear(input_dim, output_dim)
        self.k = nn.Linear(input_dim, output_dim)
        self.v = nn.Linear(input_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        
        q = self.dropout(self.q(x))
        k = self.dropout(self.k(x))
        v = self.dropout(self.v(x))

        alpha = torch.softmax(q.transpose(-2, -1) @ k, dim=-1)
        return torch.matmul(v, alpha)