欢迎来到啾啾的博客🐱。
记录学习点滴。分享工作思考和实用技巧，偶尔也分享一些杂谈💬。
有很多很多不足的地方，欢迎评论交流，感谢您的阅读和评论😄。
引言
AI使用声明：在内容整理、结构优化和语言表达的过程中，我使用了人工智能（AI）工具作为辅助。
如果以LLM应用工程师为目标，其实我们并不需要熟练掌握PyTorch，熟练掌握Transformer，但是我们必须对这两者与其背后的信息有基本的了解诶，进而更好的团队协作，以及微调模型。
本篇是一个完整的从0开始构建Transformer的Demo。
代码由QWen3-Coder生成，可以运行调试。
1 一个完整的Transformer模型
![[从零构建TransformerP2-新闻分类Demo.png]]
2 需要准备的“工具包”
| 工具 | 作用 |
|---|---|
nn.Embedding |
词嵌入 |
nn.Linear |
投影层 |
F.softmax, F.relu |
激活函数 |
torch.matmul |
矩阵乘法（注意力核心） |
mask（triu, masked_fill） |
实现因果注意力 |
LayerNorm, Dropout |
稳定训练 |
nn.ModuleList |
堆叠多层 |
DataLoader |
批量加载数据 |
3 Demo
"""
基于Transformer的新闻分类模型
严格按照设计流程实现，每个组件都有明确设计依据
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.utils.data import Dataset, DataLoader
from typing import Dict, List, Optional, Tuple
# ==============================================
# 第一部分：基础组件设计（根据设计决策选择）
# ==============================================
class TokenEmbedding(nn.Module):
"""
词嵌入层：将输入的词ID映射为密集向量表示
设计依据：
- 文本任务需要词嵌入表示语义
- 乘以sqrt(d_model)稳定初始化方差（原论文做法）
"""
def __init__(self, vocab_size: int, d_model: int):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.d_model = d_model
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
前向传播
参数:
x: 输入词ID张量，形状为(batch_size, seq_len)
返回:
嵌入后的张量，形状为(batch_size, seq_len, d_model)
""" # 原论文建议乘以sqrt(d_model)来稳定方差
return self.embedding(x) * math.sqrt(self.d_model)
class PositionalEncoding(nn.Module):
"""
位置编码：为输入序列添加位置信息
设计依据：
- Transformer没有顺序感知能力，必须添加位置信息
- 选择可学习位置编码（更灵活，适合变长序列）
"""
def __init__(self, d_model: int, max_len: int = 512):
super().__init__()
self.pos_embedding = nn.Embedding(max_len, d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
前向传播
参数:
x: 输入张量，形状为(batch_size, seq_len, d_model)
返回:
添加位置编码后的张量
""" batch_size, seq_len = x.size(0), x.size(1)
# 生成位置ID: [0, 1, 2, ..., seq_len-1]
positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
return x + self.pos_embedding(positions)
class MultiHeadAttention(nn.Module):
"""
多头注意力机制
设计依据：
- 需要建模词与词之间的关系（自注意力）
- 多头机制允许模型在不同子空间关注不同关系
"""
def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
super().__init__()
assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = 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)
self.dropout = nn.Dropout(dropout)
def scaled_dot_product_attention(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
mask: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
缩放点积注意力
参数:
q: 查询张量，形状为(batch_size, num_heads, seq_len, d_k)
k: 键张量，形状为(batch_size, num_heads, seq_len, d_k)
v: 值张量，形状为(batch_size, num_heads, seq_len, d_k)
mask: 注意力掩码，用于屏蔽padding或未来位置
返回:
attention_output: 注意力输出
attention_weights: 注意力权重（可用于可视化）
""" attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
# 将mask为0的位置设为极小值，使softmax后为0
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
attn_probs = F.softmax(attn_scores, dim=-1)
attn_probs = self.dropout(attn_probs)
output = torch.matmul(attn_probs, v)
return output, attn_probs
def split_heads(self, x: torch.Tensor) -> torch.Tensor:
"""将输入拆分为多个头"""
batch_size = x.size(0)
x = x.view(batch_size, -1, self.num_heads, self.d_k)
return x.transpose(1, 2) # (batch_size, num_heads, seq_len, d_k)
def combine_heads(self, x: torch.Tensor) -> torch.Tensor:
"""将多个头合并回原始形状"""
batch_size = x.size(0)
x = x.transpose(1, 2).contiguous()
return x.view(batch_size, -1, self.d_model)
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
mask: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""
前向传播
参数:
q, k, v: 查询、键、值张量，形状为(batch_size, seq_len, d_model)
mask: 注意力掩码
返回:
多头注意力输出，形状为(batch_size, seq_len, d_model)
""" q = self.split_heads(self.W_q(q))
k = self.split_heads(self.W_k(k))
v = self.split_heads(self.W_v(v))
attn_output, _ = self.scaled_dot_product_attention(q, k, v, mask)
output = self.W_o(self.combine_heads(attn_output))
return output
class FeedForward(nn.Module):
"""
前馈神经网络
设计依据：
- 每个位置独立处理，增强模型表示能力
- 通常d_ff = 4 * d_model（原论文比例）
"""
def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = F.gelu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)
return x
class EncoderLayer(nn.Module):
"""
编码器层
设计依据：
- 新闻分类需要双向上下文理解
- 残差连接和层归一化提升训练稳定性
"""
def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
self.ffn = FeedForward(d_model, d_ff, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
# 自注意力 + 残差连接 + 层归一化
attn_output = self.self_attn(x, x, x, mask)
x = self.norm1(x + self.dropout(attn_output))
# 前馈网络 + 残差连接 + 层归一化
ffn_output = self.ffn(x)
x = self.norm2(x + self.dropout(ffn_output))
return x
# ==============================================
# 第二部分：完整模型组装（根据设计决策）
# ==============================================
class NewsClassifier(nn.Module):
"""
新闻分类Transformer模型
设计决策回顾：
- 任务类型：文本分类（Encoder-only）
- 输入：新闻文本序列
- 输出：新闻类别（体育、科技、娱乐等）
- 架构选择：Encoder-only（无需生成能力）
- 输入表示：Token Embedding + 可学习位置编码
- 输出头：[CLS] token + 分类层
""" def __init__(
self,
vocab_size: int,
d_model: int = 768,
num_heads: int = 12,
num_layers: int = 6,
d_ff: int = 3072,
num_classes: int = 10,
max_len: int = 512,
dropout: float = 0.1
):
"""
参数:
vocab_size: 词汇表大小
d_model: 模型维度（默认768，与BERT-base一致）
num_heads: 注意力头数（默认12，与BERT-base一致）
num_layers: 编码器层数（默认6，平衡性能与计算成本）
d_ff: FFN隐藏层维度（默认3072=4*d_model）
num_classes: 分类类别数
max_len: 最大序列长度
dropout: dropout概率
""" super().__init__()
self.d_model = d_model
# 1. 特殊token（设计依据：BERT-style分类需要[CLS]）
self.cls_token = nn.Parameter(torch.randn(1, 1, d_model))
# 2. 词嵌入层
self.token_embedding = TokenEmbedding(vocab_size, d_model)
# 3. 位置编码（设计依据：选择可学习位置编码）
self.pos_encoding = PositionalEncoding(d_model, max_len)
# 4. 编码器层堆叠
self.encoder_layers = nn.ModuleList([
EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)
])
# 5. 分类头（设计依据：使用[CLS] token进行分类）
self.classifier = nn.Sequential(
nn.Linear(d_model, d_model),
nn.GELU(),
nn.Linear(d_model, num_classes)
)
self.dropout = nn.Dropout(dropout)
# 权重初始化（设计依据：稳定训练）
self._init_weights()
def _init_weights(self):
"""初始化模型权重"""
for module in self.modules():
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
nn.init.normal_(module.weight, mean=0.0, std=0.02)
elif isinstance(module, nn.LayerNorm):
nn.init.ones_(module.weight)
nn.init.zeros_(module.bias)
def add_cls_token(self, x: torch.Tensor) -> torch.Tensor:
"""
在序列开头添加[CLS] token
设计依据：BERT-style分类使用[CLS]聚合全局信息
参数:
x: 输入张量，形状为(batch_size, seq_len, d_model)
返回:
添加[CLS]后的张量，形状为(batch_size, seq_len+1, d_model)
""" batch_size = x.size(0)
cls_tokens = self.cls_token.expand(batch_size, -1, -1)
return torch.cat((cls_tokens, x), dim=1)
def create_padding_mask(self, input_ids: torch.Tensor, pad_idx: int = 0) -> torch.Tensor:
"""
创建padding掩码
设计依据：处理变长序列，忽略padding位置
参数:
input_ids: 输入ID张量，形状为(batch_size, seq_len)
pad_idx: padding token的ID
返回:
掩码张量，形状为(batch_size, 1, 1, seq_len)
True表示有效位置，False表示padding位置 (BoolTensor) """ # 创建布尔掩码，非pad为True
mask = (input_ids != pad_idx).unsqueeze(1).unsqueeze(2) # (batch_size, 1, 1, seq_len)
return mask.bool() # 确保返回的是布尔类型
def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
前向传播
参数:
input_ids: 输入词ID，形状为(batch_size, original_seq_len)
attention_mask: 可选的注意力掩码，形状为(batch_size, original_seq_len)。
1.0 表示有效位置，0.0 表示padding位置。
如果提供，应为浮点类型 (如 torch.float) 或布尔类型。
如果为 None，则根据 input_ids 自动创建。
返回:
分类logits，形状为(batch_size, num_classes)
""" batch_size, original_seq_len = input_ids.size()
# 1. 词嵌入
x = self.token_embedding(input_ids) # (batch_size, original_seq_len, d_model)
# 2. 添加[CLS] token
x = self.add_cls_token(x) # (batch_size, original_seq_len + 1, d_model)
new_seq_len = x.size(1) # 获取添加[CLS]后的序列长度
# 3. 位置编码
x = self.pos_encoding(x)
x = self.dropout(x)
# 4. 准备注意力掩码 (用于屏蔽padding)
if attention_mask is not None:
# 如果提供了 attention_mask，确保其为四维且为布尔类型
# 预期输入形状: (batch_size, original_seq_len)
# 目标形状: (batch_size, 1, 1, original_seq_len)
if attention_mask.dim() == 2:
# 假设非零值为有效位置
attention_mask_for_padding = (attention_mask != 0).unsqueeze(1).unsqueeze(2)
elif attention_mask.dim() == 4:
attention_mask_for_padding = (attention_mask.squeeze(1).squeeze(1) != 0).unsqueeze(1).unsqueeze(2)
else:
raise ValueError(f"attention_mask must be 2D or 4D, but got {attention_mask.dim()}D")
else:
# 如果没有提供，根据 input_ids 自动创建
# 形状: (batch_size, 1, 1, original_seq_len)
attention_mask_for_padding = self.create_padding_mask(input_ids)
# --- 关键修复：正确扩展 mask 以适应添加了 [CLS] token 后的新序列长度 --- # 创建一个针对新序列长度 (new_seq_len = original_seq_len + 1) 的掩码
# [CLS] token (索引 0) 应该总是被 attend 到，所以我们需要扩展 mask # 1. 初始化一个全为 True 的新掩码，形状 (batch_size, 1, 1, new_seq_len) expanded_mask = torch.ones((batch_size, 1, 1, new_seq_len), dtype=torch.bool, device=x.device)
# 2. 将原始 padding mask 复制到新 mask 的 [1:] 位置 (跳过 [CLS]) # 原始 mask 形状: (batch_size, 1, 1, original_seq_len)
# 新 mask 的 [1:] 部分形状: (batch_size, 1, 1, original_seq_len)
expanded_mask[:, :, :, 1:] = attention_mask_for_padding
# 最终用于注意力的掩码，形状 (batch_size, 1, 1, new_seq_len) # 在 MultiHeadAttention 中，这个掩码会被广播用于屏蔽 key (src_seq) 的 padding 位置
final_attention_mask = expanded_mask
# 5. 通过编码器层
# 将扩展后的 mask 传递给每一层，以屏蔽 padding for layer in self.encoder_layers:
x = layer(x, final_attention_mask) # 传递匹配新序列长度的 mask
# 6. 取[CLS] token作为句子表示
cls_output = x[:, 0, :] # (batch_size, d_model)
# 7. 分类
logits = self.classifier(cls_output)
return logits
# ==============================================
# 第三部分：训练流程（根据设计决策）
# ==============================================
def train_news_classifier():
"""新闻分类模型训练流程"""
# 1. 超参数设置（根据设计决策）
config = {
"vocab_size": 30000, # 词汇表大小（设计依据：新闻领域常用词）
"d_model": 768, # 模型维度（设计依据：平衡性能与计算成本）
"num_heads": 12, # 注意力头数（设计依据：与d_model匹配）
"num_layers": 6, # 编码器层数（设计依据：足够捕捉复杂关系）
"d_ff": 3072, # FFN维度（设计依据：4*d_model）
"num_classes": 10, # 分类类别数（设计依据：新闻类别数量）
"max_len": 512, # 最大序列长度（设计依据：覆盖大多数新闻）
"dropout": 0.1, # dropout概率（设计依据：防止过拟合）
"batch_size": 32, # 批量大小（设计依据：GPU内存限制）
"learning_rate": 2e-5, # 学习率（设计依据：微调预训练模型常用值）
"epochs": 3, # 训练轮数（设计依据：避免过拟合）
"warmup_steps": 500, # warmup步数（设计依据：稳定训练初期）
"weight_decay": 0.01 # 权重衰减（设计依据：正则化）
}
# 2. 创建模型
print("✅ 创建新闻分类模型...")
model = NewsClassifier(
vocab_size=config["vocab_size"],
d_model=config["d_model"],
num_heads=config["num_heads"],
num_layers=config["num_layers"],
d_ff=config["d_ff"],
num_classes=config["num_classes"],
max_len=config["max_len"],
dropout=config["dropout"]
)
# 3. 设备选择
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f" 模型将运行在: {device}")
# 4. 伪造数据集（实际应用中替换为真实数据）
class NewsDataset(Dataset):
def __init__(self, num_samples: int = 1000, max_len: int = 512):
self.num_samples = num_samples
self.max_len = max_len
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
# 伪造新闻文本（词ID）
seq_len = min(500, 100 + idx % 400) # 变长序列
input_ids = torch.randint(1, 30000, (seq_len,))
# 伪造类别标签（0-9）
label = torch.tensor(idx % 10, dtype=torch.long)
return input_ids, label
# 5. 数据加载器（处理变长序列的关键）
def collate_fn(batch):
"""处理变长序列的collate函数"""
input_ids, labels = zip(*batch)
# 找出最大长度
max_len = max(len(ids) for ids in input_ids)
# padding
padded_ids = []
for ids in input_ids:
padding = torch.zeros(max_len - len(ids), dtype=torch.long)
padded_ids.append(torch.cat([ids, padding]))
input_ids = torch.stack(padded_ids)
labels = torch.stack(labels)
return input_ids, labels
print("✅ 创建数据集和数据加载器...")
train_dataset = NewsDataset(num_samples=1000)
train_loader = DataLoader(
train_dataset,
batch_size=config["batch_size"],
shuffle=True,
collate_fn=collate_fn
)
# 6. 损失函数和优化器
print("✅ 配置训练组件...")
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(
model.parameters(),
lr=config["learning_rate"],
weight_decay=config["weight_decay"]
)
# 7. 学习率调度器（设计依据：warmup + linear decay）
total_steps = len(train_loader) * config["epochs"]
warmup_steps = config["warmup_steps"]
def lr_lambda(current_step: int):
if current_step < warmup_steps:
return float(current_step) / float(max(1, warmup_steps))
return max(
0.0, float(total_steps - current_step) / float(max(1, total_steps - warmup_steps))
)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
# 8. 训练循环
print("🚀 开始训练...")
for epoch in range(config["epochs"]):
model.train()
total_loss = 0
for batch_idx, (input_ids, labels) in enumerate(train_loader):
input_ids = input_ids.to(device)
labels = labels.to(device)
# 前向传播
optimizer.zero_grad()
logits = model(input_ids)
loss = loss_fn(logits, labels)
# 反向传播
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 梯度裁剪
optimizer.step()
scheduler.step()
total_loss += loss.item()
# 打印进度
if batch_idx % 50 == 0:
avg_loss = total_loss / (batch_idx + 1)
current_lr = optimizer.param_groups[0]['lr']
print(f"Epoch [{epoch + 1}/{config['epochs']}] | "
f"Batch [{batch_idx}/{len(train_loader)}] | "
f"Loss: {avg_loss:.4f} | " f"LR: {current_lr:.2e}")
print(f"✅ Epoch {epoch + 1} 完成 | Average Loss: {total_loss / len(train_loader):.4f}")
# 9. 保存模型
torch.save(model.state_dict(), "news_classifier.pth")
print("💾 模型已保存至 news_classifier.pth")
# ==============================================
# 第四部分：推理示例
# ==============================================
def predict_news_category(text: str, model: NewsClassifier, tokenizer, device: torch.device):
"""
新闻分类推理
设计依据：
- 使用与训练相同的预处理流程
- 取[CLS] token进行分类
参数:
text: 新闻文本
model: 训练好的模型
tokenizer: 文本分词器
device: 设备
返回:
预测类别和概率
""" model.eval()
# 1. 文本预处理
input_ids = tokenizer.encode(text, max_length=512, truncation=True, padding="max_length")
input_ids = torch.tensor(input_ids).unsqueeze(0).to(device)
# 2. 前向传播
with torch.no_grad():
logits = model(input_ids)
probs = F.softmax(logits, dim=-1)
# 3. 获取结果
predicted_class = torch.argmax(probs, dim=-1).item()
confidence = probs[0, predicted_class].item()
return predicted_class, confidence
if __name__ == "__main__":
# 这里只是演示结构，实际运行需要完整实现
print("=" * 50)
print("Transformer新闻分类模型设计与实现")
print("=" * 50)
print("\n本示例演示了如何根据任务需求设计并实现一个Transformer模型")
print("设计流程严格遵循：问题分析 → 架构选择 → 组件设计 → 训练实现")
print("\n关键设计决策：")
print("- 选择Encoder-only架构（分类任务无需生成能力）")
print("- 使用[CLS] token进行分类（BERT-style）")
print("- 可学习位置编码（更适合变长新闻文本）")
print("- 6层编码器（平衡性能与计算成本）")
print("\n要运行完整训练，请取消注释train_news_classifier()调用")
train_news_classifier()