手写attention

#!/user/bin/env python3
# -*- coding: utf-8 -*-
# @Time     : 2022-09-27 20:30
# @Author   : Lyt
# @IDE      : PyCharm    
# @FileName : Attention.py
# @Blog     : https://blog.csdn.net/m0_53292725?type=blog
import torch.nn as nn
import torch
from einops import rearrange


class Attention(nn.Module):
    def __init__(self, dim, dim_head=64, heads=8, dropout=0.):
        super(Attention, self).__init__()

        inner_dim = dim_head * heads
        self.heads = heads
        self.to_qkv = nn.Linear(dim, inner_dim*3)
        self.softmax = nn.Softmax(dim=-1)
        self.scale = dim_head ** -0.5
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
        dots = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        atten = self.softmax(dots)
        out = torch.matmul(atten, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)