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

1次阅读

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

AG News 数据集作为文本分类的经典基准数据集，包含 4 个新闻类别（World、Sports、Business、Sci/Tech）共 12 万条样本。但在实际应用中常遇到以下问题：

HTML 标签残留 ：原始数据包含大量<br> 等 HTML 标签，直接用正则处理可能遗漏特殊符号
类别分布不均：Sports 类样本明显多于 Sci/Tech 类（约 1.3:1 比例）
长文本处理：标题 + 正文拼接后平均长度达 210 词，远超 BERT 的 512token 限制
噪声数据：存在非 ASCII 字符和编码错误样本

传统正则表达式 re.sub(r'<[^>]+>', '', text) 可能无法处理嵌套标签，建议组合使用 BeautifulSoup：

from bs4 import BeautifulSoup
def clean_html(raw_text):
    soup = BeautifulSoup(raw_text, "html.parser")
    return soup.get_text(separator=" ")

额外处理步骤：

统一全半角字符
过滤非打印字符
处理 URL 占位符（如 [URL] 标记）

import re
def normalize_text(text):
    text = re.sub(r'\[.*?\]', '', text)  # 移除方括号内容
    text = re.sub(r'\s+', ' ', text)      # 合并空白符
    return text.strip()

适用于轻量级场景：

优势：训练速度快，内存消耗低
劣势：无法捕获语义关联

from sklearn.feature_extraction.text import TfidfVectorizer
tfidf = TfidfVectorizer(max_features=5000, ngram_range=(1,2))
X_train = tfidf.fit_transform(train_texts)

推荐 HuggingFace 实现：

使用 bert-base-uncased 分词器
对长文本采用滑动窗口策略

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
inputs = tokenizer(
    text,
    truncation=True,
    padding='max_length', 
    max_length=256,
    return_tensors="pt"
)

模型类型	参数量	训练时间	准确率
BiLSTM+Attention	4.7M	23min	89.2%
CNN-Text	3.1M	18min	90.1%
DistilBERT	66M	47min	93.7%

使用生成器避免内存爆炸：

class NewsDataset(Dataset):
    def __init__(self, texts, labels):
        self.texts = texts
        self.labels = labels

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

    def __getitem__(self, idx):
        # 动态生成样本，不预加载全部数据
        text = self.texts[idx]
        label = self.labels[idx]
        inputs = tokenizer(text, **tokenizer_kwargs)
        return {**inputs, 'labels': torch.tensor(label)}

显著减少显存占用：

from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()

with autocast():
    outputs = model(**batch)
    loss = outputs.loss
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

所有预处理步骤（如 TF-IDF 拟合）必须仅在训练集进行
使用 sklearn.model_selection.train_test_split 时设置固定 random_state

验证动态轴设置是否正确
检查操作符支持情况（如 F.grid_sample 需要 opset>=11）

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input_ids", "attention_mask"],
    dynamic_axes={"input_ids": {0: "batch"},
        "attention_mask": {0: "batch"}
    }
)