谷歌transformer模型简单介绍

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

# ==================== 总说明 ====================
# 本代码实现了一个简化版Transformer的单编码器层，包含：
# 1. 输入嵌入 + 位置编码
# 2. 多头自注意力机制
# 3. 前馈网络
# 4. 残差连接 + 层归一化
# 注意：为简化代码，省略了完整的解码器、掩码机制和训练流程
# ===================================================

# -------------------- 超参数定义 --------------------
d_model = 512       # 模型维度（词向量维度）
n_heads = 8         # 注意力头数
d_ff = 2048         # 前馈网络隐藏层维度
dropout = 0.1       # Dropout概率
max_len = 100       # 最大序列长度
vocab_size = 10000  # 词汇表大小

# -------------------- 输入嵌入层 --------------------
class Embeddings(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.lut = nn.Embedding(vocab_size, d_model)  # 词嵌入矩阵
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * torch.sqrt(torch.tensor(self.d_model))  # 缩放词向量

# -------------------- 位置编码 --------------------
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)  # 初始化位置编码矩阵
        position = torch.arange(0, max_len).unsqueeze(1)  # 生成位置索引 [max_len,1]
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model))  # 计算缩放系数
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维使用sin函数
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维使用cos函数
        pe = pe.unsqueeze(0)  # 增加batch维度 [1,max_len,d_model]
        self.register_buffer('pe', pe)  # 注册为不可训练参数

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]  # 将位置编码加到输入上
        return self.dropout(x)

# -------------------- 多头注意力 --------------------
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        self.d_k = d_model // n_heads  # 每个头的维度
        self.n_heads = n_heads
        self.linears = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(4)])  # 4个线性层（Q/K/V/输出）
        self.dropout = nn.Dropout(p=dropout)

    def scaled_dot_product_attention(self, Q, K, V):
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k))  # 缩放点积
        attn = F.softmax(scores, dim=-1)  # Softmax归一化
        return torch.matmul(attn, V), attn  # 返回加权和与注意力权重

    def split_heads(self, x):
        batch_size = x.size(0)
        return x.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # 分割多头 [batch,heads,seq_len,d_k]

    def forward(self, Q, K, V):
        Q, K, V = [self.linears[i](x) for i, x in enumerate([Q, K, V])]  # 线性变换
        Q, K, V = [self.split_heads(x) for x in [Q, K, V]]  # 分割多头
        x, attn = self.scaled_dot_product_attention(Q, K, V)  # 计算注意力
        x = x.transpose(1, 2).contiguous().view(x.size(0), -1, self.d_model)  # 合并多头
        return self.linears[-1](x)  # 最终线性变换

# -------------------- 前馈网络 --------------------
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)  # 第一层线性变换
        self.w2 = nn.Linear(d_ff, d_model)  # 第二层线性变换
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w2(self.dropout(F.relu(self.w1(x))))  # ReLU激活 + Dropout

# -------------------- 编码器层 --------------------
class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)  # 多头注意力
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)  # 前馈网络
        self.norm1 = nn.LayerNorm(d_model)  # 层归一化1
        self.norm2 = nn.LayerNorm(d_model)  # 层归一化2
        self.dropout1 = nn.Dropout(dropout)  # Dropout1
        self.dropout2 = nn.Dropout(dropout)  # Dropout2

    def forward(self, x):
        # 自注意力子层（含残差连接）
        attn_out = self.self_attn(x, x, x)  # Q=K=V
        x = self.norm1(x + self.dropout1(attn_out))  # 残差连接 + 层归一化
        # 前馈网络子层（含残差连接）
        ff_out = self.feed_forward(x)
        x = self.norm2(x + self.dropout2(ff_out))  # 残差连接 + 层归一化
        return x

# -------------------- 测试代码 --------------------
if __name__ == "__main__":
    batch_size, seq_len = 32, 10
    x = torch.randint(0, vocab_size, (batch_size, seq_len))  # 随机生成输入序列
    embed = Embeddings(vocab_size, d_model)  # 初始化嵌入层
    pos_enc = PositionalEncoding(d_model, dropout, max_len)  # 初始化位置编码
    encoder_layer = EncoderLayer(d_model, n_heads, d_ff, dropout)  # 初始化编码器层
    x_embed = embed(x)  # 词嵌入 [batch,seq_len,d_model]
    x_pos = pos_enc(x_embed)  # 加入位置编码
    output = encoder_layer(x_pos)  # 编码器层处理 [batch,seq_len,d_model]
    print("输出形状:", output.shape)  # 验证输出维度