从零实现文本分类:PyTorch搭建TextCNN中文分类器实战
2026/9/8 10:25:21 网站建设 项目流程

简介:这套基于PyTorch的TextCNN中文文本分类项目,提供完整代码与可直接运行的实验数据,面向NLP初学者及需要快速落地中文分类任务的开发者。压缩包共35个文件、约73.2MB,其中13个Python脚本覆盖数据加载、分词、模型定义、训练与评估全流程,另有文本数据样本、PB模型权重、说明文档与效果截图,目录结构清晰,便于按模块对照学习。资源在CSDN上已有近800人学习参考,实践反馈良好。项目不仅详细展示了词嵌入、卷积、最大池化、全连接等TextCNN核心模块的PyTorch实现,还针对中文分词、序列填充、类别不平衡等预处理难点给出了配套处理思路。通过运行该工程,读者能掌握从数据清洗到模型评估的完整链路,并可方便地迁移至情感分析、主题分类、新闻分类等常见NLP场景。 中文NLP入门绕不开一个最经典的需求:文本分类。做新闻归档、评论情感分析、工单自动打标,本质上都是让模型判断“这段文字属于哪个类别”。而在刚起步的阶段,TextCNN 几乎是绕不过去的模型,用 PyTorch 实现它代码量小、逻辑清晰,特别适合拿来当第一个能完整跑通的中文文本分类项目。这篇文章从数据和代码角度把事情讲透,围绕一个可以直接运行的 TextCNN PyTorch 工程展开,你既可以照着敲一遍,也可以把它当作脚手架,快速改造成你自己的分类任务。

1. TextCNN 为什么值得作为中文分类的起步模型

1.1 模型核心思想:不同尺寸卷积核就是不同类型的 n-gram 探测器

TextCNN 的思路并不复杂:先把句子里的每个词映射成一个向量,得到“词向量矩阵”,然后用多个宽度不同的卷积核在这个矩阵上做扫描。每个卷积核只看连续的几个词,宽度为 3 的卷积核捕获“三连词”特征,宽度为 5 的捕获“五连词”特征,最后通过池化层把最关键的特征挑出来,送到全连接层做分类。

你可以把它想成一组“滑动放大镜”:小窗口负责看局部搭配,比如“非常棒”“太好用”;大窗口负责看稍微长一点的语义块,比如“物流速度非常快”。不同尺寸的卷积核各看各的,最后汇总,这比单纯判断某几个词是否出现要灵活很多。对于中文短文本分类来说,这种 n-gram 级别的特征提取方式非常有效,因为很多分类信号本来就藏在这类局部搭配里。

1.2 它和 BERT、BiLSTM、传统方法的位置差异

很多初学者一上来就想上 BERT,但如果是小数据集、CPU 环境、或者只是想快速出一个可用的基线,BERT 的性价比并不高。我把常见方案的取舍整理成了表格,方便你判断:

模型训练成本数据需求可解释性适合场景
TextCNN中低较好短文本、基线模型、快速验证
BiLSTM中高一般序列依赖明显的任务
BERT较差数据量大、GPU 充足、追求极限精度
朴素贝叶斯等传统方法最低简单规则、冷启动

TextCNN 最大的价值是“稳”:结构不会出错,训练不容易崩,收敛速度也比循环神经网络快。用它跑通整个工程链路之后,再换 BERT 或者 BiLSTM,你会发现大部分代码骨架都能复用,真正要改的只有模型定义部分。

2. 环境准备与工程目录设计

2.1 依赖清单与 PyTorch 安装

这个项目依赖很少,核心只需要 Python 3.8 以上、PyTorch、NumPy 和小部分数据处理工具。PyTorch 安装没有想象中的复杂,CPU 环境直接执行:

pip install torch numpy pandas jieba scikit-learn

如果你本机有 NVIDIA GPU,想用 GPU 加速,建议去 PyTorch 官网选择合适的 CUDA 版本安装命令,不要随便从镜像仓库拉老版本。判断 PyTorch 是否能用 GPU,用下面这一行最直接:

import torch print(torch.__version__) print(torch.cuda.is_available())

个人建议:第一次跑这个工程完全可以用 CPU,数据量小的时候两者差别不明显,反而能省掉很多驱动层面的麻烦。

2.2 一个清晰的工程目录

工程文件不要随便扔在一起。我建议按下面这个结构组织:

textcnn_project/ ├── data/ │ └── news_sample.csv ├── src/ │ ├── preprocess.py │ ├── models.py │ ├── train.py │ └── predict.py └── checkpoints/

data 目录放原始数据,src 目录放处理逻辑、模型定义、训练脚本和预测脚本,checkpoints 目录留给模型权重。这样划分之后,后续换数据、换模型、调参数都不会牵一发而动全身。

3. 数据处理:从 CSV 到可训练的样本

3.1 数据集格式与读取方式

这里我准备的是一个中文短文本新闻标题数据集,两列格式:label 是类别,text 是标题内容。示例数据长这样:

label,text 科技,三款国产旗舰手机横向评测出炉 体育,中超联赛第十四轮比赛今晚开打 财经,央行开展千亿元逆回购操作 娱乐,电影暑期档票房突破五十亿 健康,秋冬季节如何预防流感 教育,多所高校公布本科招生计划

读取和查看数据分布只需要几行代码:

import pandas as pd df = pd.read_csv("data/news_sample.csv") print(df["label"].value_counts())

需要注意两点。第一,不要用中文做标签值,尽量换成拼音或英文如 tech、sports,否则后续做标签到数字的映射时容易出编码问题。第二,每类数据量尽量保持均匀,如果某类只有 5 条而其他类有 500 条,模型基本学不到这个少数类的特征。

3.2 中文分词、词表构建与序列化

中文文本不能像英文那样按空格切词,需要先分词。这里使用 jieba 做正向最大匹配式的分词,简单粗暴但够用:

import jieba def cut_text(text: str): return [w for w in jieba.cut(text) if w.strip()]

切完词之后要做两件事:构建词表、把词序列转成索引序列。词表构建时我会给未登录词预留一个<unk>位置,也给填充位留一个<pad>。注意 pad 这个位置很关键,后面模型 Embedding 层要用 padding_idx 指定它,避免 pad 符号参与梯度更新。

PAD_TOKEN = "<pad>" UNK_TOKEN = "<unk>" def build_vocab(texts, min_count=1): word_count = {} for text in texts: for word in cut_text(text): word_count[word] = word_count.get(word, 0) + 1 words = [PAD_TOKEN, UNK_TOKEN] + [w for w, c in word_count.items() if c >= min_count] word2idx = {w: idx for idx, w in enumerate(words)} return word2idx

句子长度不可能完全一致,所以还需要一个编码函数,把每个句子统一截断或填充到固定长度,比如 64:

def encode(text, word2idx, max_len=64): ids = [word2idx.get(w, word2idx[UNK_TOKEN]) for w in cut_text(text)] if len(ids) > max_len: ids = ids[:max_len] else: ids = ids + [word2idx[PAD_TOKEN]] * (max_len - len(ids)) return ids

3.3 Dataset 与 DataLoader 封装

PyTorch 里最规范的做法是继承 Dataset 并实现__len____getitem__。还有一个容易踩坑的地方是 collate_fn,尤其是 batch 中每条样本长度不一致时,必须在 collate_fn 里做 pad 操作。我这里提前统一成了 max_len,所以 DataLoader 不会报维度错:

from torch.utils.data import Dataset, DataLoader class TextDataset(Dataset): def __init__(self, df, word2idx, max_len=64): self.df = df.reset_index(drop=True) self.word2idx = word2idx self.max_len = max_len self.labels = sorted(df["label"].unique()) self.label2idx = {label: i for i, label in enumerate(self.labels)} def __len__(self): return len(self.df) def __getitem__(self, index): row = self.df.iloc[index] ids = encode(row["text"], self.word2idx, self.max_len) label = self.label2idx[row["label"]] return torch.tensor(ids, dtype=torch.long), torch.tensor(label, dtype=torch.long)

4. TextCNN 模型完整实现

4.1 模型定义:Embedding + 多尺寸卷积 + 池化 + 全连接

下面这段是我常用的 TextCNN 结构。代码不长,但每个模块都有明确作用。我会保留注释,方便你逐行对照卷积核维度变化:

import torch import torch.nn as nn import torch.nn.functional as F class TextCNN(nn.Module): def __init__(self, vocab_size, embed_dim=128, num_filters=128, filter_sizes=(3, 4, 5), num_classes=6, dropout=0.5, pad_idx=0): super().__init__() self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx) self.convs = nn.ModuleList([ nn.Conv2d(1, num_filters, (kernel_size, embed_dim), padding=(kernel_size // 2, 0)) for kernel_size in filter_sizes ]) self.fc = nn.Linear(len(filter_sizes) * num_filters, num_classes) self.dropout = nn.Dropout(dropout) def forward(self, x): # x shape: [batch_size, seq_len] emb = self.embedding(x) # [batch, seq_len, embed_dim] emb = emb.unsqueeze(1) # [batch, 1, seq_len, embed_dim] conv_outputs = [] for conv in self.convs: c = conv(emb) # [batch, num_filters, seq_len, 1] c = c.squeeze(3) # 去掉最后一维 pooled = F.max_pool1d(c, c.size(2)).squeeze(2) # [batch, num_filters] conv_outputs.append(pooled) feature = torch.cat(conv_outputs, dim=1) # 拼接不同卷积核提取的特征 feature = self.dropout(feature) logits = self.fc(feature) return logits

这里需要理解维度变化。Embedding 输出是三维张量,Conv2d 要求输入是四维,所以先加一个维度和通道维类似的维度。卷积核形状是(kernel_size, embed_dim),在序列长度和词向量维度两个方向上滑动,每次卷积后输出通道是 num_filters。max_pool1d会把每个卷积核在整条序列上找到的最大值提取出来,这个操作叫 max-pooling over time,它保证无论句子多长,最后进入全连接层的特征长度都是固定的。

4.2 损失函数、优化器与训练主循环

分类任务损失函数选交叉熵,优化器我习惯用 Adam,学习率从 1e-3 开始,效果不行再降。训练主循环要注意切分训练集和验证集,每一轮都记录 loss 和准确率,方便肉眼观察有没有过拟合。

import torch.optim as optim from sklearn.model_selection import train_test_split def train_model(model, train_loader, val_loader, epochs=10, lr=1e-3): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=lr) for epoch in range(epochs): model.train() total_loss, correct, total = 0, 0, 0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() logits = model(inputs) loss = criterion(logits, labels) loss.backward() optimizer.step() total_loss += loss.item() * inputs.size(0) correct += (logits.argmax(1) == labels).sum().item() total += labels.size(0) train_acc = correct / total val_acc = evaluate(model, val_loader, device) print(f"epoch {epoch+1}/{epochs} loss={total_loss/total:.4f} train_acc={train_acc:.4f} val_acc={val_acc:.4f}")

注意训练前一定要设置model.train(),评估前要设置model.eval()。dropout 在训练和推理阶段行为不一样,忘了切换会导致预测结果波动很大。

5. 模型评估、保存与单条预测

5.1 验证集评估与混淆矩阵

评估函数不需要改动模型,只做前向传播。为了防止梯度计算浪费显存和内存,要包在torch.no_grad()里:

def evaluate(model, val_loader, device): model.eval() correct, total = 0, 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) logits = model(inputs) correct += (logits.argmax(1) == labels).sum().item() total += labels.size(0) return correct / total

如果数据类别不均衡,只看准确率会骗人,还要看混淆矩阵。可以用 sklearn 快速生成,看模型到底把哪些类别互相混在一起:

from sklearn.metrics import confusion_matrix, classification_report model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for inputs, labels in val_loader: inputs = inputs.to(device) logits = model(inputs) all_preds.extend(logits.argmax(1).cpu().numpy()) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds))

5.2 模型保存、加载与单条预测

保存模型时建议只保存 state_dict,不保存整个模型对象,这样模型结构升级后依旧能加载旧权重:

torch.save(model.state_dict(), "checkpoints/textcnn.pt")

加载时先实例化一个结构完全相同的模型,再把状态字典填进去:

model = TextCNN( vocab_size=len(word2idx), num_classes=len(label2idx) ) model.load_state_dict(torch.load("checkpoints/textcnn.pt", map_location="cpu")) model.eval()

单条预测的完整流程是这样的:分词、查词表、编码成固定长度序列、构造 batch 维度、模型推理、取概率最大的类别索引,最后映射回标签名。

def predict(text, model, word2idx, label2idx, max_len=64): idx_to_label = {v: k for k, v in label2idx.items()} ids = encode(text, word2idx, max_len) input_tensor = torch.tensor([ids], dtype=torch.long) model.eval() with torch.no_grad(): logits = model(input_tensor) pred_idx = logits.argmax(1).item() return idx_to_label[pred_idx]

这里要注意 encode 函数必须和训练时保持一致,尤其是 max_len 和分词逻辑。如果你训练时用的是 jieba,预测时却用普通字符切分,结果基本不可用。

6. 实操中的踩坑记录与调参建议

6.1 高频报错定位:维度问题与词表缺失

我跑这个工程时遇到过几个典型问题,写出来给你排雷。

第一个是卷积层报“Expected 4D input”。原因一般是忘记对 Embedding 输出做unsqueeze(1)。Conv2d 只接受四维张量,形状应该是[batch, 1, seq_len, embed_dim]

第二个是词表缺失。有些数据里出现的词在词表里没有,导致运行时 KeyError。这种情况要么是构建词表和训练集不一致,要么是词频过滤阈值太高。解决方式是在 encode 函数里统一用word2idx.get(w, word2idx[UNK_TOKEN]),保证所有词都有兜底路径。

第三个是 DataLoader 在读取不同长度样本时报错。最省事的办法是像我上面那样,在数据集内部提前 pad 到固定长度;如果数据长短差异特别大,则应该在 collate_fn 里用pad_sequence,同时保留 attention_mask 之类的辅助信息,但这个工程用不到,统一 max_len 就足够。

6.2 让效果更稳的四个建议

第一,固定随机种子。PyTorch、Python random、NumPy 三处都要固定,否则每次跑结果不同,很难判断改动是否有效。

import random import numpy as np def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)

第二,用预训练词向量初始化 Embedding。如果词表不大,随机初始化效果也可接受;但如果数据量不够,建议加载公开的中文词向量或者直接换成 BERT 做特征抽取。

第三,类别不均衡时给 CrossEntropyLoss 传 class weights。这样少数类样本虽然数量少,但梯度贡献会被放大,不至于完全被多数类淹没。

第四,早停。我想特别强调一下:别一门心思等几十个 epoch 跑完,在第 10 轮之后如果验证集准确率连续几轮不涨甚至下降,直接停掉,保存验证集准确率最高的那一次权重即可。

这个工程我前前后后跑了两遍,第一次在 CPU 上用小型标题数据,几秒一个 epoch;第二次换到 GPU,反而要处理 cuda 环境、num_workers 并发这些杂事。个人体会是:新手最好先用最小数据集把前向传播、反向传播、评估预测这条链路完整跑通,再考虑上规模。之后无论是换数据集、加预训练词向量,还是把模型整体升级成更复杂的结构,这套工程骨架都能够继续用。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询