当前位置：首页 > article >正文

从零开始：使用PyTorch构建DeepSeek R1模型及其训练详解

article 2025/2/26 19:36:11

本文将引导你使用 PyTorch 从零开始构建 DeepSeek R1 模型，并详细解释模型架构和训练步骤。DeepSeek R1 是一个假设的模型名称，为了演示目的，我们将构建一个基于 Transformer 的简单文本生成模型。

1. 模型架构

DeepSeek R1 的核心是一个基于 Transformer 的编码器-解码器架构，包含以下关键组件：

Embedding Layer: 将输入的单词索引转换为密集向量表示。
Positional Encoding: 为输入序列添加位置信息，因为 Transformer 本身不具备处理序列顺序的能力。
Encoder: 由多个编码器层堆叠而成，每个编码器层包含：
- Multi-Head Self-Attention: 捕捉输入序列中不同位置之间的依赖关系。
- Feed-Forward Network: 对每个位置的表示进行非线性变换。
Decoder: 由多个解码器层堆叠而成，每个解码器层包含：
- Masked Multi-Head Self-Attention: 防止解码器在预测下一个单词时看到未来的信息。
- Multi-Head Encoder-Decoder Attention: 允许解码器关注编码器的输出。
- Feed-Forward Network: 对每个位置的表示进行非线性变换。
Output Layer: 将解码器的输出转换为词汇表上的概率分布。

2. 代码实现

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

class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, max_seq_length, dropout=0.1):
        super(Transformer, self).__init__()
        
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = nn.Parameter(torch.zeros(1, max_seq_length, d_model))
        
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers)
        
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers)
        
        self.fc_out = nn.Linear(d_model, vocab_size)
        
    def forward(self, src, tgt):
        src_seq_length, tgt_seq_length = src.size(1), tgt.size(1)
        
        src = self.embedding(src) + self.positional_encoding[:, :src_seq_length, :]
        tgt = self.embedding(tgt) + self.positional_encoding[:, :tgt_seq_length, :]
        
        memory = self.encoder(src)
        output = self.decoder(tgt, memory)
        
        return self.fc_out(output)

# 定义超参数
vocab_size = 10000
d_model = 512
nhead = 8
num_encoder_layers = 6
num_decoder_layers = 6
dim_feedforward = 2048
max_seq_length = 100
dropout = 0.1

# 初始化模型
model = Transformer(vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, max_seq_length, dropout)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)

# 训练循环
for epoch in range(10):
    for src, tgt in dataloader:
        optimizer.zero_grad()
        
        output = model(src, tgt[:, :-1])
        loss = criterion(output.reshape(-1, vocab_size), tgt[:, 1:].reshape(-1))
        
        loss.backward()
        optimizer.step()
        
    print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

3. 分步训练详解

数据准备: 将文本数据转换为模型可接受的格式，例如将单词映射到索引，并将数据分批。
模型初始化: 使用定义的超参数初始化模型。
损失函数和优化器: 选择交叉熵损失函数和 Adam 优化器。
训练循环:
- 将输入序列 (src) 和目标序列 (tgt) 输入模型。
- 模型输出预测的下一个单词的概率分布。
- 计算预测分布和目标序列之间的损失。
- 反向传播损失并更新模型参数。
评估: 使用验证集评估模型性能，例如计算困惑度 (perplexity)。