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

评论363

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

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. import math
  5.  
  6. # 定义位置编码器
  7. class PositionalEncoder(nn.Module):
  8.     def __init__(self, d_model, max_len=512):
  9.         super(PositionalEncoder, self).__init__()
  10.         self.d_model = d_model
  11.         self.max_len = max_len
  12.         
  13.         # 计算位置编码表
  14.         pe = torch.zeros(max_len, d_model)  # 创建一个形状为(max_len, d_model)的零张量pe，用于存储位置编码
  15.         position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # 创建一个形状为(max_len, 1)的张量，表示位置
  16.         div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))  # 计算位置编码中的分母部分
  17.         pe[:, 0::2] = torch.sin(position * div_term)  # 计算位置编码中偶数位置的值
  18.         pe[:, 1::2] = torch.cos(position * div_term)  # 计算位置编码中奇数位置的值
  19.         pe = pe.unsqueeze(0)  # 在第0维上增加一维，用于处理批次数据
  20.         self.register_buffer('pe', pe)  # 将位置编码表pe注册为模型的缓冲区
  21.         
  22.     def forward(self, x):
  23.         # 输入x的维度为(batch_size, seq_len, d_model)
  24.         x = x * math.sqrt(self.d_model)  # 对输入乘以一个缩放因子，以便缓解梯度消失问题
  25.         seq_len = x.size(1)  # 获取输入序列的长度
  26.         # 将位置编码添加到输入中
  27.         x = x + self.pe[:, :seq_len]  # 在对应位置添加位置编码
  28.         return x
  29.  
  30. # 定义多头注意力机制
  31. class MultiHeadAttention(nn.Module):
  32.     def __init__(self, d_model, num_heads):
  33.         super(MultiHeadAttention, self).__init__()
  34.         assert d_model % num_heads == 0, "d_model必须被num_heads整除"
  35.         self.d_model = d_model
  36.         self.num_heads = num_heads
  37.         self.head_dim = d_model // num_heads  # 每个头的维度
  38.         
  39.         # 定义线性变换层
  40.         self.W_q = nn.Linear(d_model, d_model)  # 查询向量的线性变换层
  41.         self.W_k = nn.Linear(d_model, d_model)  # 键向量的线性变换层
  42.         self.W_v = nn.Linear(d_model, d_model)  # 值向量的线性变换层
  43.         self.W_o = nn.Linear(d_model, d_model)  # 输出向量的线性变换层
  44.         
  45.     def forward(self, query, key, value, mask=None):
  46.         batch_size = query.size(0)  # 获取批次大小
  47.         
  48.         # 将输入的query、key、value通过线性变换得到Q、K、V
  49.         Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # 计算查询向量Q
  50.         K = self.W_k(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # 计算键向量K
  51.         V = self.W_v(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # 计算值向量V
  52.         
  53.         # 计算注意力分数
  54.         attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)  # 计算注意力分数
  55.         if mask is not None:
  56.             attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))  # 使用mask处理注意力分数
  57.         attention_weights = F.softmax(attention_scores, dim=-1)  # 计算注意力权重
  58.         
  59.         # 计算注意力值
  60.         attention_output = torch.matmul(attention_weights, V)  # 计算注意力值
  61.         attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)  # 调整注意力值的形状
  62.         
  63.         # 经过线性变换得到最终输出
  64.         output = self.W_o(attention_output)  # 最终输出
  65.         return output
  66.  
  67. # 定义前向传播层
  68. class FeedForward(nn.Module):
  69.     def __init__(self, d_model, d_ff):
  70.         super(FeedForward, self).__init__()
  71.         self.d_model = d_model
  72.         self.d_ff = d_ff
  73.         
  74.         # 定义两个线性变换层
  75.         self.linear1 = nn.Linear(d_model, d_ff)  # 第一个线性变换层
  76.         self.linear2 = nn.Linear(d_ff, d_model)  # 第二个线性变换层
  77.         
  78.     def forward(self, x):
  79.         x = F.relu(self.linear1(x))  # 使用ReLU激活函数进行非线性变换
  80.         x = self.linear2(x)  # 进行第二个线性变换
  81.         return x
  82.  
  83. # 定义一个Transformer模型
  84. class Transformer(nn.Module):
  85.     def __init__(self, d_model, num_heads, d_ff, num_layers):
  86.         super(Transformer, self).__init__()
  87.         self.d_model = d_model
  88.         self.num_heads = num_heads
  89.         self.d_ff = d_ff
  90.         self.num_layers = num_layers
  91.         
  92.         # 定义多个编码器层
  93.         self.encoder_layers = nn.ModuleList([
  94.             nn.ModuleList([
  95.                 MultiHeadAttention(d_model, num_heads),
  96.                 nn.LayerNorm(d_model),
  97.                 FeedForward(d_model, d_ff),
  98.                 nn.LayerNorm(d_model)
  99.             ])
  100.             for _ in range(num_layers)
  101.         ])
  102.         
  103.     def forward(self, src, mask=None):
  104.         x = src
  105.         
  106.         # 通过多个编码器层进行前向传播
  107.         for i in range(self.num_layers):
  108.             # 多头注意力层
  109.             attention = self.encoder_layers[i][0]
  110.             norm1 = self.encoder_layers[i][1]
  111.             x = x + attention(x, x, x, mask=mask)
  112.             x = norm1(x)
  113.             
  114.             # 前向传播层
  115.             feed_forward = self.encoder_layers[i][2]
  116.             norm2 = self.encoder_layers[i][3]
  117.             x = x + feed_forward(x)
  118.             x = norm2(x)
  119.         
  120.         return x
  121.  
  122. # 测试Transformer模型
  123. if __name__ == "__main__":
  124.     # 假设输入维度为(16, 20, 512)，即(batch_size, seq_len, d_model)
  125.     input_tensor = torch.randn(16, 20, 512)
  126.     transformer_model = Transformer(d_model=512, num_heads=8, d_ff=2048, num_layers=6)
  127.     output = transformer_model(input_tensor)
  128.     print(output.shape)  # 输出：torch.Size([16, 20, 512])

 

以下是longfromer的pytorch实现版本

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. import math
  5.  
  6. # 定义位置编码器
  7. class PositionalEncoder(nn.Module):
  8.     def __init__(self, d_model, max_len=512):
  9.         super(PositionalEncoder, self).__init__()
  10.         self.d_model = d_model
  11.         self.max_len = max_len
  12.         
  13.         # 计算位置编码表
  14.         pe = torch.zeros(max_len, d_model)
  15.         position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
  16.         div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
  17.         pe[:, 0::2] = torch.sin(position * div_term)
  18.         pe[:, 1::2] = torch.cos(position * div_term)
  19.         pe = pe.unsqueeze(0)
  20.         self.register_buffer('pe', pe)
  21.         
  22.     def forward(self, x):
  23.         x = x * math.sqrt(self.d_model)
  24.         seq_len = x.size(1)
  25.         x = x + self.pe[:, :seq_len]
  26.         return x
  27.  
  28. # 定义局部注意力机制
  29. class LocalAttention(nn.Module):
  30.     def __init__(self, d_model, local_window):
  31.         super(LocalAttention, self).__init__()
  32.         self.d_model = d_model
  33.         self.local_window = local_window
  34.         self.attention = nn.MultiheadAttention(d_model, 1)
  35.         
  36.     def forward(self, x, mask=None):
  37.         batch_size, seq_len, _ = x.size()
  38.         local_mask = torch.zeros(seq_len, seq_len)
  39.         for i in range(seq_len):
  40.             local_mask[i, max(0, i - self.local_window):i + self.local_window + 1] = 1
  41.         local_mask = local_mask.unsqueeze(0).to(x.device)
  42.         local_mask = local_mask * mask.unsqueeze(1) if mask is not None else local_mask
  43.         
  44.         return self.attention(x.permute(1, 0, 2), x.permute(1, 0, 2), x.permute(1, 0, 2), key_padding_mask=local_mask)
  45.  
  46. # 定义Longformer模型
  47. class Longformer(nn.Module):
  48.     def __init__(self, d_model, num_heads, d_ff, num_layers, local_window):
  49.         super(Longformer, self).__init__()
  50.         self.d_model = d_model
  51.         self.num_heads = num_heads
  52.         self.d_ff = d_ff
  53.         self.num_layers = num_layers
  54.         self.local_window = local_window
  55.         
  56.         # 定义位置编码器和局部注意力层
  57.         self.positional_encoder = PositionalEncoder(d_model)
  58.         self.local_attention = LocalAttention(d_model, local_window)
  59.         
  60.         # 定义多个编码器层
  61.         self.encoder_layers = nn.ModuleList([
  62.             nn.ModuleList([
  63.                 nn.LayerNorm(d_model),
  64.                 nn.Linear(d_model, d_ff),
  65.                 nn.ReLU(),
  66.                 nn.Linear(d_ff, d_model),
  67.                 nn.LayerNorm(d_model)
  68.             ])
  69.             for _ in range(num_layers)
  70.         ])
  71.         
  72.     def forward(self, src, mask=None):
  73.         x = self.positional_encoder(src)
  74.         for i in range(self.num_layers):
  75.             norm1 = self.encoder_layers[i][0]
  76.             linear1 = self.encoder_layers[i][1]
  77.             relu = self.encoder_layers[i][2]
  78.             linear2 = self.encoder_layers[i][3]
  79.             norm2 = self.encoder_layers[i][4]
  80.             
  81.             # 局部注意力层
  82.             if mask is not None:
  83.                 mask[:, :, :self.local_window] = 0
  84.             x = x + self.local_attention(x.permute(1, 0, 2), mask=mask)[0].permute(1, 0, 2)
  85.             
  86.             # 前向传播层
  87.             x = norm1(x)
  88.             x = linear2(relu(linear1(x))) + x
  89.             x = norm2(x)
  90.         
  91.         return x
  92.  
  93. # 测试Longformer模型
  94. if __name__ == "__main__":
  95.     input_tensor = torch.randn(16, 512, 512)  # 假设输入维度为(16, 512, 512)
  96.     mask = torch.ones(16, 512)  # 假设有512个标记
  97.     longformer_model = Longformer(d_model=512, num_heads=8, d_ff=2048, num_layers=6, local_window=128)
  98.     output = longformer_model(input_tensor, mask=mask)
  99.     print(output.shape)  # 输出：torch.Size([16, 512, 512])

 

发表评论

匿名网友

拖动滑块以完成验证
加载失败