PyTorch基于Transformer架构的完整文本生成实现方案

pip install torch torchvision torchaudio transformers datasets accelerate

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer
import numpy as np
from tqdm import tqdm

# ==================== 配置参数 ====================
class Config:
    vocab_size = 50257  # GPT-2 tokenizer 词汇表大小
    d_model = 768       # 模型维度
    nhead = 12          # 注意力头数
    num_layers = 12     # Transformer 层数
    dropout = 0.1
    batch_size = 8
    seq_len = 128       # 序列长度
    learning_rate = 3e-4
    num_epochs = 10
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

config = Config()

# ==================== 数据集 ====================
class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=128):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.encodings = []
        
        for text in texts:
            encoding = tokenizer(
                text,
                truncation=True,
                padding='max_length',
                max_length=max_length,
                return_tensors='pt'
            )
            self.encodings.append({
                'input_ids': encoding['input_ids'].squeeze(),
                'attention_mask': encoding['attention_mask'].squeeze()
            })
    
    def __len__(self):
        return len(self.encodings)
    
    def __getitem__(self, idx):
        return self.encodings[idx]

# ==================== Transformer 模型 ====================
class TransformerLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # 词嵌入层
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_embedding = nn.Embedding(config.seq_len, config.d_model)
        
        # Transformer 编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.d_model * 4,
            dropout=config.dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, config.num_layers)
        
        # 输出层
        self.fc_out = nn.Linear(config.d_model, config.vocab_size)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, x, mask=None):
        # 位置编码
        batch_size, seq_len = x.shape
        positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0)
        
        # 嵌入 + 位置编码
        x = self.embedding(x) + self.pos_embedding(positions)
        x = self.dropout(x)
        
        # Transformer 编码
        transformer_out = self.transformer(x, src_key_padding_mask=~mask.bool() if mask is not None else None)
        
        # 输出预测
        output = self.fc_out(transformer_out)
        return output

# ==================== 训练函数 ====================
def train_model(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    
    progress_bar = tqdm(dataloader, desc="Training")
    for batch in progress_bar:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        
        # 创建目标（右移一位）
        targets = input_ids[:, 1:].contiguous()
        input_ids = input_ids[:, :-1].contiguous()
        attention_mask = attention_mask[:, :-1].contiguous()
        
        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask)
        
        # 计算损失（忽略填充位置）
        loss = criterion(outputs.view(-1, config.vocab_size), targets.view(-1))
        loss.backward()
        
        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        total_loss += loss.item()
        progress_bar.set_postfix({'loss': loss.item()})
    
    return total_loss / len(dataloader)

# ==================== 文本生成函数 ====================
def generate_text(model, tokenizer, prompt, max_length=50, temperature=1.0, top_k=50):
    model.eval()
    with torch.no_grad():
        # 编码输入提示
        input_ids = tokenizer.encode(prompt, return_tensors='pt').to(config.device)
        generated = input_ids
        
        for _ in range(max_length):
            # 获取模型输出
            outputs = model(generated)
            next_token_logits = outputs[:, -1, :] / temperature
            
            # Top-k 采样
            if top_k > 0:
                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
                next_token_logits[indices_to_remove] = -float('Inf')
            
            # Softmax + 采样
            probs = torch.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            # 检查是否生成结束符
            if next_token.item() == tokenizer.eos_token_id:
                break
                
            generated = torch.cat([generated, next_token], dim=-1)
        
        return tokenizer.decode(generated[0], skip_special_tokens=True)

# ==================== 主训练流程 ====================
def main():
    # 初始化 tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
    tokenizer.pad_token = tokenizer.eos_token
    
    # 准备示例数据（实际使用时替换为真实数据集）
    sample_texts = [
        "Artificial intelligence is transforming the world.",
        "Machine learning models require large amounts of data.",
        "Natural language processing enables computers to understand human language.",
        "Deep learning has achieved remarkable success in various domains.",
        "Transformer architecture revolutionized sequence modeling."
    ] * 100  # 重复以创建足够数据
    
    # 创建数据集和数据加载器
    dataset = TextDataset(sample_texts, tokenizer, config.seq_len)
    dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)
    
    # 初始化模型
    model = TransformerLM(config).to(config.device)
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
    optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate)
    
    # 训练循环
    print(f"Starting training on {config.device}...")
    for epoch in range(config.num_epochs):
        avg_loss = train_model(model, dataloader, optimizer, criterion, config.device)
        print(f"Epoch {epoch+1}/{config.num_epochs}, Average Loss: {avg_loss:.4f}")
        
        # 每 2 个 epoch 生成示例文本
        if (epoch + 1) % 2 == 0:
            prompt = "Artificial intelligence"
            generated_text = generate_text(model, tokenizer, prompt, max_length=30)
            print(f"Generated text: {generated_text}\n")
    
    # 保存模型
    torch.save(model.state_dict(), 'transformer_lm.pth')
    print("Model saved successfully!")

if __name__ == "__main__":
    main()

# 绝对位置编码 vs 相对位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        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() * (-np.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):
        return x + self.pe[:, :x.size(1)]

# 多头注意力计算过程
def scaled_dot_product_attention(q, k, v, mask=None):
    """
    q, k, v: [batch_size, seq_len, d_k]
    """
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / np.sqrt(d_k)  # [B, L, L]
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attn_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v)
    return output, attn_weights

# 防止梯度爆炸
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 预热 + 余弦退火
def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * (current_step - num_warmup_steps) / 
                                               float(max(1, num_training_steps - num_warmup_steps)))))
    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

# 忽略填充 token 的损失计算
criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

# 控制生成多样性
next_token_logits = outputs[:, -1, :] / temperature
# temperature < 1: 更确定性
# temperature > 1: 更随机性

# Top-k 采样
def top_k_sampling(logits, k):
    indices_to_remove = logits < torch.topk(logits, k)[0][..., -1, None]
    logits[indices_to_remove] = -float('Inf')
    return logits

# Top-p (Nucleus) 采样
def top_p_sampling(logits, p):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_remove = cumulative_probs > p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0
    indices_to_remove = sorted_indices_to_remove.scatter(
        dim=-1, index=sorted_indices, src=sorted_indices_to_remove
    )
    logits[indices_to_remove] = -float('Inf')
    return logits

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        outputs = model(input_ids)
        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# 多 GPU 训练
model = nn.DataParallel(model)
# 或使用 DistributedDataParallel (更高效)
model = nn.parallel.DistributedDataParallel(model)

# 动态量化
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments

# 加载预训练模型
model = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token

# 训练参数
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
)

# 训练器
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

trainer.train()

from transformers import pipeline

# 使用 pipeline 进行文本生成
generator = pipeline('text-generation', model='gpt2', device=0)

result = generator(
    "Artificial intelligence is",
    max_length=50,
    num_return_sequences=1,
    temperature=0.7,
    top_k=50,
    top_p=0.95
)
print(result[0]['generated_text'])