基于pytorch实现transformer以及longformer以及代码详解

以下是一个基于PyTorch实现Transformer模型的简单示例代码，并对每个步骤进行了详细的注释。请注意，这个示例主要着重于模型的实现，不包括数据处理和训练部分。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# 定义位置编码器
class PositionalEncoder(nn.Module):
    def __init__(self, d_model, max_len=512):
        super(PositionalEncoder, self).__init__()
        self.d_model = d_model
        self.max_len = max_len
        
        # 计算位置编码表
        pe = torch.zeros(max_len, d_model)  # 创建一个形状为(max_len, d_model)的零张量pe，用于存储位置编码
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # 创建一个形状为(max_len, 1)的张量，表示位置
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))  # 计算位置编码中的分母部分
        pe[:, 0::2] = torch.sin(position * div_term)  # 计算位置编码中偶数位置的值
        pe[:, 1::2] = torch.cos(position * div_term)  # 计算位置编码中奇数位置的值
        pe = pe.unsqueeze(0)  # 在第0维上增加一维，用于处理批次数据
        self.register_buffer('pe', pe)  # 将位置编码表pe注册为模型的缓冲区
        
    def forward(self, x):
        # 输入x的维度为(batch_size, seq_len, d_model)
        x = x * math.sqrt(self.d_model)  # 对输入乘以一个缩放因子，以便缓解梯度消失问题
        seq_len = x.size(1)  # 获取输入序列的长度
        # 将位置编码添加到输入中
        x = x + self.pe[:, :seq_len]  # 在对应位置添加位置编码
        return x

# 定义多头注意力机制
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model必须被num_heads整除"
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads  # 每个头的维度
        
        # 定义线性变换层
        self.W_q = nn.Linear(d_model, d_model)  # 查询向量的线性变换层
        self.W_k = nn.Linear(d_model, d_model)  # 键向量的线性变换层
        self.W_v = nn.Linear(d_model, d_model)  # 值向量的线性变换层
        self.W_o = nn.Linear(d_model, d_model)  # 输出向量的线性变换层
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)  # 获取批次大小
        
        # 将输入的query、key、value通过线性变换得到Q、K、V
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # 计算查询向量Q
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # 计算键向量K
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # 计算值向量V
        
        # 计算注意力分数
        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)  # 计算注意力分数
        if mask is not None:
            attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))  # 使用mask处理注意力分数
        attention_weights = F.softmax(attention_scores, dim=-1)  # 计算注意力权重
        
        # 计算注意力值
        attention_output = torch.matmul(attention_weights, V)  # 计算注意力值
        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)  # 调整注意力值的形状
        
        # 经过线性变换得到最终输出
        output = self.W_o(attention_output)  # 最终输出
        return output

# 定义前向传播层
class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(FeedForward, self).__init__()
        self.d_model = d_model
        self.d_ff = d_ff
        
        # 定义两个线性变换层
        self.linear1 = nn.Linear(d_model, d_ff)  # 第一个线性变换层
        self.linear2 = nn.Linear(d_ff, d_model)  # 第二个线性变换层
        
    def forward(self, x):
        x = F.relu(self.linear1(x))  # 使用ReLU激活函数进行非线性变换
        x = self.linear2(x)  # 进行第二个线性变换
        return x

# 定义一个Transformer模型
class Transformer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers):
        super(Transformer, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        self.num_layers = num_layers
        
        # 定义多个编码器层
        self.encoder_layers = nn.ModuleList([
            nn.ModuleList([
                MultiHeadAttention(d_model, num_heads),
                nn.LayerNorm(d_model),
                FeedForward(d_model, d_ff),
                nn.LayerNorm(d_model)
            ])
            for _ in range(num_layers)
        ])
        
    def forward(self, src, mask=None):
        x = src
        
        # 通过多个编码器层进行前向传播
        for i in range(self.num_layers):
            # 多头注意力层
            attention = self.encoder_layers[i][0]
            norm1 = self.encoder_layers[i][1]
            x = x + attention(x, x, x, mask=mask)
            x = norm1(x)
            
            # 前向传播层
            feed_forward = self.encoder_layers[i][2]
            norm2 = self.encoder_layers[i][3]
            x = x + feed_forward(x)
            x = norm2(x)
        
        return x

# 测试Transformer模型
if __name__ == "__main__":
    # 假设输入维度为(16, 20, 512)，即(batch_size, seq_len, d_model)
    input_tensor = torch.randn(16, 20, 512)
    transformer_model = Transformer(d_model=512, num_heads=8, d_ff=2048, num_layers=6)
    output = transformer_model(input_tensor)
    print(output.shape)  # 输出：torch.Size([16, 20, 512])

以下是longfromer的pytorch实现版本

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# 定义位置编码器
class PositionalEncoder(nn.Module):
    def __init__(self, d_model, max_len=512):
        super(PositionalEncoder, self).__init__()
        self.d_model = d_model
        self.max_len = max_len
        
        # 计算位置编码表
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x * math.sqrt(self.d_model)
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len]
        return x

# 定义局部注意力机制
class LocalAttention(nn.Module):
    def __init__(self, d_model, local_window):
        super(LocalAttention, self).__init__()
        self.d_model = d_model
        self.local_window = local_window
        self.attention = nn.MultiheadAttention(d_model, 1)
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        local_mask = torch.zeros(seq_len, seq_len)
        for i in range(seq_len):
            local_mask[i, max(0, i - self.local_window):i + self.local_window + 1] = 1
        local_mask = local_mask.unsqueeze(0).to(x.device)
        local_mask = local_mask * mask.unsqueeze(1) if mask is not None else local_mask
        
        return self.attention(x.permute(1, 0, 2), x.permute(1, 0, 2), x.permute(1, 0, 2), key_padding_mask=local_mask)

# 定义Longformer模型
class Longformer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers, local_window):
        super(Longformer, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        self.num_layers = num_layers
        self.local_window = local_window
        
        # 定义位置编码器和局部注意力层
        self.positional_encoder = PositionalEncoder(d_model)
        self.local_attention = LocalAttention(d_model, local_window)
        
        # 定义多个编码器层
        self.encoder_layers = nn.ModuleList([
            nn.ModuleList([
                nn.LayerNorm(d_model),
                nn.Linear(d_model, d_ff),
                nn.ReLU(),
                nn.Linear(d_ff, d_model),
                nn.LayerNorm(d_model)
            ])
            for _ in range(num_layers)
        ])
        
    def forward(self, src, mask=None):
        x = self.positional_encoder(src)
        for i in range(self.num_layers):
            norm1 = self.encoder_layers[i][0]
            linear1 = self.encoder_layers[i][1]
            relu = self.encoder_layers[i][2]
            linear2 = self.encoder_layers[i][3]
            norm2 = self.encoder_layers[i][4]
            
            # 局部注意力层
            if mask is not None:
                mask[:, :, :self.local_window] = 0
            x = x + self.local_attention(x.permute(1, 0, 2), mask=mask)[0].permute(1, 0, 2)
            
            # 前向传播层
            x = norm1(x)
            x = linear2(relu(linear1(x))) + x
            x = norm2(x)
        
        return x

# 测试Longformer模型
if __name__ == "__main__":
    input_tensor = torch.randn(16, 512, 512)  # 假设输入维度为(16, 512, 512)
    mask = torch.ones(16, 512)  # 假设有512个标记
    longformer_model = Longformer(d_model=512, num_heads=8, d_ff=2048, num_layers=6, local_window=128)
    output = longformer_model(input_tensor, mask=mask)
    print(output.shape)  # 输出：torch.Size([16, 512, 512])