code样例 - 技术栈

SelfAttention

import torch

import torch.nn as nn

from math import sqrt

torch.bmm 专用于批量矩阵乘法，适用于形状为 (batch_size, n, m) 和 (batch_size, m, p) 的 3D 张量。

torch.matmul 支持更灵活的张量乘法运算

class SelfAttention(nn.Module):

def init (self, input_dim, dim_k, dim_v):

self.q = nn.Linear(input_dim, dim_k)

self.k = nn.Linear(input_dim, dim_k)

self.v = nn.Linear(input_dim, dim_v)

复制代码

    self.norm = np.sqrt(dim_k)

def forward(self, x):
    Q = self.q(x) # [bs, seq_len, dim_k]
    K = self.k(x)
    V = self.v(x)
    atten = nn.Softmax(dim=-1)(torch.bmm(Q, K.permute(0, 2, 1)) / self.norm)
    return torch.bmm(atten, V)


##CrossAttention
import torch

import torch.nn as nn

查询通常来自解码器，键和值通常来自编码器

class CrossAttention(nn.Module):

def init (self, input_dim, dim_k, dim_v):

self.q = nn.Linear(input_dim, dim_k)

self.k = nn.Linear(input_dim, dim_k)

self.v = nn.Linear(input_dim, dim_v)

复制代码

    self.norm = np.sqrt(dim_k)

def forward(self, x1, x2):
    Q1 = self.q(x1) # [bs, seq_len, dim_k]
    K2 = self.k(x2)
    V2 = self.v(x2)
    atten = nn.Softmax(dim=-1)(torch.bmm(Q1, K2.permute(0, 2, 1)) / self.norm)
    return torch.bmm(atten, V2)

## CausalAttention
class CausalAttention(nn.Module):
def __init__(self, input_dim, dim_k, dim_v):
    super(CausalAttention, self).__init__()
    self.q = nn.Linear(input_dim, dim_k)
    self.k = nn.Linear(input_dim, dim_k)
    self.v = nn.Linear(input_dim, dim_v)
    self.norm = np.sqrt(dim_k)

def forward(self, x):
    # Linear projections for query, key, and value
    Q = self.q(x)  # [batch_size, seq_len, dim_k]
    K = self.k(x)
    V = self.v(x)
    
    atten = torch.bmm(Q, K.transpose(1, 2)) / self.norm  # [batch_size, seq_len, seq_len]
    
    seq_len = atten.size(-1)
    mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)
    atten = atten.masked_fill(mask == 0, float('-inf'))
    
    mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)
    atten = atten.masked_fill(mask == 0, float('-inf'))
    
    atten = nn.Softmax(dim=-1)(atten)
    return torch.bmm(atten, V)

# 多头注意力机制

class MultiHeadSelfAttention(nn.Module):

def init (self, input_dim, heads, dim_k, dim_v):

super(MultiHeadSelfAttention, self).init ()

self.heads = heads

self.dim_k = dim_k // heads

self.dim_v = dim_v // heads

复制代码

    self.q = nn.Linear(input_dim, dim_k)
    self.k = nn.Linear(input_dim, dim_k)
    self.v = nn.Linear(input_dim, dim_v)

    self._norm_fact = 1 / sqrt(self.dim_k)
    self.out = nn.Linear(self.dim_v, input_dim)

def forward(self, x):
    bs = x.size(0)

    Q = self.q(x).view(bs, -1, self.heads, self.dim_k).permute(0, 2, 1, 3)
    K = self.k(x).view(bs, -1, self.heads, self.dim_k).permute(0, 2, 1, 3)
    V = self.v(x).view(bs, -1, self.heads, self.dim_v).permute(0, 2, 1, 3)

    atten = nn.Softmax(dim=-1)(torch.matmul(Q, K.permute(0, 1, 3, 2))) * self._norm_fact
    atten = torch.matmul(atten, V)

    atten = atten.permute(0, 2, 1, 3).contiguous().view(bs, -1, self.dim_v)
    return self.out(atten)  


# 位置编码

class PositionalEncoding(nn.Module):

def init (self, d_model, max_len=5000):

super(PositionalEncoding, self).init ()

self.encoding = torch.zeros(max_len, d_model)

position = torch.arange(0, max_len).unsqueeze(1).float()

div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))

self.encoding $:, 0::2$ = torch.sin(position * div_term)

self.encoding $:, 1::2$ = torch.cos(position * div_term)

self.encoding = self.encoding.unsqueeze(0)

复制代码

def forward(self, x):
    return self.encoding[:, :x.size(1)]

生成上三角矩阵

def get_attn_subsequence_mask(seq): # seq: $batch_size, tgt_len$

attn_shape = $seq.size(0), seq.size(1), seq.size(1)$ # 生成上三角矩阵, $batch_size, tgt_len, tgt_len$

subsequence_mask = np.triu(np.ones(attn_shape), k=1)

subsequence_mask = torch.from_numpy(subsequence_mask).byte() # $batch_size, tgt_len, tgt_len$

return subsequence_mask