AG News数据集实战指南：从数据预处理到文本分类模型构建

1次阅读

共计 2074 个字符，预计需要花费 6 分钟才能阅读完成。

AG News 是新闻文本分类的经典数据集，包含 4 个类别（World、Sports、Business、Sci/Tech）共 12 万条数据。原始数据存在几个典型问题：

格式混乱 ：文本中混杂 HTML 标签（如 <br>）和特殊字符（如®™）
长度不均 ：新闻正文从几十字到上千字不等，直接输入模型会导致内存溢出
类别偏差 ：各类别样本量差异达 15%（Sports 类最多，Sci/Tech 类最少）

数据清洗 ：正则表达式去噪 + NLTK 停用词过滤
文本向量化 ：对比 TF-IDF（适合传统机器学习）和 Word2Vec（适合深度学习）
模型选型 ：采用 TextCNN（计算效率高，适合短文本）

词嵌入维度：Word2Vec 选择 300 维（平衡效果与计算成本）
文本截断：统一处理为 512 字符（覆盖 90% 样本）
损失函数：加权交叉熵（权重与类别数量成反比）

import re
from nltk.corpus import stopwords

def clean_text(text):
    # 移除 HTML 标签
    text = re.sub(r'<[^>]+>', '', text) 
    # 保留字母、数字和基本标点
    text = re.sub(r'[^\w\s.,!?]', '', text)
    # 转换为小写并分词
    words = text.lower().split()
    # 过滤停用词
    stops = set(stopwords.words('english'))
    return [w for w in words if w not in stops]

from torch.utils.data import Dataset, DataLoader

class AGNewsDataset(Dataset):
    def __init__(self, texts, labels, max_len=512):
        self.texts = [t[:max_len] for t in texts]  # 截断长文本
        self.labels = labels

    def __getitem__(self, idx):
        return self.texts[idx], self.labels[idx]

    def __len__(self):
        return len(self.texts)

# 示例用法
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

import torch.nn as nn

class TextCNN(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)

        # 3 种不同尺度的卷积核
        self.convs = nn.ModuleList([nn.Conv2d(1, 100, (k, embed_dim)) 
            for k in [3,4,5]
        ])
        self.fc = nn.Linear(300, 4)  # 4 分类

    def forward(self, x):
        x = self.embedding(x)  # [batch, seq, embed]
        x = x.unsqueeze(1)  # 增加通道维

        # 多尺度卷积 + 池化
        features = [nn.functional.relu(conv(x)).squeeze(3)  # [batch, 100, seq-k+1]
            for conv in self.convs
        ]
        pooled = [nn.functional.max_pool1d(f, f.size(2)).squeeze(2) 
            for f in features
        ]

        return self.fc(torch.cat(pooled, 1))

动态截断 ：保留前 512 字符（覆盖 90% 样本）
关键信息优先 ：实验发现新闻关键信息多出现在首段

# 计算类别权重
class_counts = [8000, 12000, 9000, 11000]  # 示例数据
weights = 1. / torch.tensor(class_counts, dtype=torch.float)
criterion = nn.CrossEntropyLoss(weight=weights)

from sklearn.model_selection import train_test_split

# 分层抽样保持类别比例
X_train, X_val, y_train, y_val = train_test_split(
    texts, labels, 
    test_size=0.2, 
    stratify=labels
)