用PyTorch和TorchText搞定AG_NEWS新闻分类：从数据加载到75%准确率的保姆级代码

用PyTorch和TorchText实现AG_NEWS新闻分类从零到75%准确率的完整指南当你第一次接触文本分类任务时可能会被数据处理和模型构建的复杂性吓到。本文将带你用PyTorch和TorchText从零开始构建一个新闻分类器无需任何先验知识只需跟着步骤操作就能达到75%的准确率。1. 环境准备与数据加载在开始之前确保你已经安装了最新版本的PyTorch和TorchText。可以通过以下命令安装pip install torch torchtextAG_NEWS是学术界常用的新闻分类数据集包含四个类别世界新闻、体育、商业和科技。让我们先加载这个数据集import torch import torch.nn as nn from torchtext.datasets import AG_NEWS from torchtext.data.utils import get_tokenizer from collections import Counter, OrderedDict # 创建数据目录并加载数据集 train_dataset, test_dataset AG_NEWS(root./data, split(train, test)) classes [World, Sports, Business, Sci/Tech]提示如果直接从GitHub下载数据集遇到问题可以尝试设置代理或更换网络环境。查看数据集前几个样本可以帮助我们理解数据结构for i, (label, text) in zip(range(3), train_dataset): print(f{classes[label-1]}: {text[:50]}...)2. 文本预处理与词表构建文本数据不能直接输入模型需要转换为数值形式。我们将使用以下步骤处理文本分词将句子拆分为单词或子词单元构建词表创建单词到索引的映射序列化将文本转换为数字序列# 使用基础英语分词器 tokenizer get_tokenizer(basic_english) # 构建词表 counter Counter() for (label, text) in train_dataset: counter.update(tokenizer(text)) # 按词频排序并创建词表 vocab torchtext.vocab.vocab( OrderedDict(sorted(counter.items(), keylambda x: x[1], reverseTrue)), min_freq1 ) vocab_size len(vocab) print(f词表大小: {vocab_size})3. 数据批处理与填充文本长度不一致是常见问题我们需要统一长度以便批量处理def padify(batch): # 将文本转换为索引序列 texts [vocab.lookup_indices(tokenizer(item[1])) for item in batch] # 获取当前批次最大长度 max_len max(map(len, texts)) # 对短文本进行填充 padded [torch.nn.functional.pad( torch.tensor(text), (0, max_len - len(text)), modeconstant, value0 ) for text in texts] labels torch.LongTensor([item[0]-1 for item in batch]) texts torch.stack(padded) return labels, texts # 创建数据加载器 train_loader torch.utils.data.DataLoader( list(train_dataset), # 转换为列表 batch_size32, collate_fnpadify, shuffleTrue )4. 模型构建与训练我们将使用简单的嵌入层全连接层的架构class NewsClassifier(nn.Module): def __init__(self, vocab_size, embed_dim, num_classes): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim) self.fc nn.Linear(embed_dim, num_classes) def forward(self, x): # 获取词嵌入 embedded self.embedding(x) # [batch_size, seq_len, embed_dim] # 平均池化 pooled torch.mean(embedded, dim1) # [batch_size, embed_dim] # 分类 return self.fc(pooled)训练函数实现def train_model(model, dataloader, epochs3, lr0.001): optimizer torch.optim.Adam(model.parameters(), lrlr) criterion nn.CrossEntropyLoss() model.train() for epoch in range(epochs): total_loss 0 correct 0 count 0 for labels, texts in dataloader: optimizer.zero_grad() outputs model(texts) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() count len(labels) print(fEpoch {epoch1}: Loss{total_loss/count:.4f}, Acc{correct/count:.2%})初始化并训练模型# 初始化模型 model NewsClassifier(vocab_size, 64, len(classes)) # 训练模型 train_model(model, train_loader, epochs3, lr0.001)5. 模型评估与优化训练完成后我们需要评估模型在测试集上的表现def evaluate(model, dataloader): model.eval() correct 0 total 0 with torch.no_grad(): for labels, texts in dataloader: outputs model(texts) correct (outputs.argmax(1) labels).sum().item() total len(labels) return correct / total test_loader torch.utils.data.DataLoader( list(test_dataset), batch_size32, collate_fnpadify ) accuracy evaluate(model, test_loader) print(f测试集准确率: {accuracy:.2%})要进一步提升性能可以考虑以下优化策略调整嵌入维度尝试32、64、128等不同维度使用预训练词向量如GloVe或FastText添加更多层在嵌入层后加入LSTM或Transformer层调整学习率尝试不同的学习率和学习率调度策略# 使用预训练词向量的示例 pretrained_embeds torchtext.vocab.GloVe(name6B, dim100) # 替换模型中的嵌入层 model.embedding.weight.data.copy_(pretrained_embeds.get_vecs_by_tokens(vocab.get_itos()))6. 实际应用与部署训练好的模型可以保存并用于实际预测# 保存模型 torch.save(model.state_dict(), news_classifier.pth) # 加载模型 loaded_model NewsClassifier(vocab_size, 64, len(classes)) loaded_model.load_state_dict(torch.load(news_classifier.pth)) # 预测函数 def predict(text, model, vocab, tokenizer): tokens tokenizer(text) indices vocab.lookup_indices(tokens) tensor torch.LongTensor(indices).unsqueeze(0) model.eval() with torch.no_grad(): output model(tensor) prob torch.softmax(output, dim1) pred_class classes[output.argmax().item()] return pred_class, prob.max().item() # 示例预测 sample_text Apple releases new iPhone with advanced AI features pred_class, confidence predict(sample_text, loaded_model, vocab, tokenizer) print(f预测类别: {pred_class}, 置信度: {confidence:.2%})7. 常见问题与解决方案在实际应用中可能会遇到以下问题内存不足减小批量大小使用更小的嵌入维度考虑使用更高效的优化器如Adagrad过拟合添加Dropout层使用L2正则化增加训练数据性能瓶颈使用DataLoader的num_workers参数并行加载数据考虑使用混合精度训练在GPU上训练# 添加Dropout的改进模型 class ImprovedNewsClassifier(nn.Module): def __init__(self, vocab_size, embed_dim, num_classes): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim) self.dropout nn.Dropout(0.5) self.fc nn.Linear(embed_dim, num_classes) def forward(self, x): embedded self.embedding(x) pooled torch.mean(embedded, dim1) pooled self.dropout(pooled) return self.fc(pooled)在项目实践中我发现以下几个技巧特别有用使用学习率调度器可以显著提升模型性能早停法(early stopping)能有效防止过拟合模型集成可以进一步提升准确率1-2个百分点