本文记录一次完整的文本分类模型训练过程。目标是训练一个模型，把输入文本分成四类：
- 前端
- 后端
- AI
- 其他
例如：
React hooks 怎么使用 -> 前端
MySQL 慢查询怎么优化 -> 后端
什么是大模型微调 -> AI
我今天晚上吃什么 -> 其他
整个过程分成两个版本：
- 先用
TF-IDF + LogisticRegression跑通完整流程 - 再升级到
Transformer / BERT版本
一、为什么先不用 Transformer？
一开始没有直接上 Transformer，而是先用简单模型，是为了先理解文本分类的完整流程。
文本分类的主流程其实是固定的：
准备数据
↓
读取文本和标签
↓
把文本转成模型能处理的数字
↓
训练模型
↓
测试模型
↓
保存模型
↓
加载模型
↓
预测新文本
TF-IDF 版本和 Transformer 版本的区别主要在中间两步。
| 简单版本 | Transformer 版本 |
|---|---|
TfidfVectorizer |
AutoTokenizer |
LogisticRegression |
BERT + 分类头 |
model.fit() |
PyTorch 训练循环 |
joblib.dump() |
save_pretrained() |
joblib.load() |
from_pretrained() |
本质都是：
文字 → 数字 → 模型 → 分类结果
第一部分：TF-IDF 文本分类模型
1. 准备数据
先准备一个 data.csv 文件：
text,label
React 组件如何传递 props,前端
Vue3 的组合式 API 怎么使用,前端
CSS flex 布局怎么居中,前端
JavaScript 事件冒泡是什么,前端
前端页面如何做响应式布局,前端
浏览器渲染机制是什么,前端
Webpack 打包优化怎么做,前端
TypeScript 泛型怎么用,前端
HTML 表单如何提交,前端
前端跨域问题怎么解决,前端
Spring Boot 怎么连接 MySQL,后端
Java 后端接口如何设计,后端
Redis 缓存穿透怎么解决,后端
Node.js 如何写 REST API,后端
数据库索引为什么能提高查询速度,后端
Nginx 如何做反向代理,后端
如何设计用户登录接口,后端
消息队列 Kafka 有什么用,后端
JWT 登录认证怎么实现,后端
Docker 如何部署后端服务,后端
机器学习模型怎么训练,AI
大模型微调 LoRA 是什么,AI
神经网络反向传播原理,AI
Transformer 注意力机制是什么,AI
如何用 PyTorch 训练分类模型,AI
Embedding 向量是什么意思,AI
BERT 模型怎么做文本分类,AI
深度学习中的损失函数是什么,AI
梯度下降算法是什么,AI
卷积神经网络 CNN 是什么,AI
今天晚上吃什么,其他
周末去哪里玩,其他
这个电影好不好看,其他
天气不错适合散步,其他
我想买一双运动鞋,其他
这个手机壳多少钱,其他
明天要不要去超市,其他
最近有什么好看的电视剧,其他
怎么学英语比较好,其他
这家餐厅好吃吗,其他
这个数据集包含两列：
text = 文本内容
label = 文本类别
2. 安装依赖
pip install pandas scikit-learn jieba joblib
3. 编写训练代码
创建 train.py：
import jieba
import joblib
import pandas as pd
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
# =========================
# 1. 读取 CSV 数据
# =========================
df = pd.read_csv("data.csv")
texts = df["text"].tolist()
labels = df["label"].tolist()
# =========================
# 2. 检查数据
# =========================
print("数据总数：", len(texts))
print("类别分布：")
print(df["label"].value_counts())
# =========================
# 3. 中文分词函数
# =========================
def chinese_tokenizer(text):
return list(jieba.cut(text))
# =========================
# 4. 划分训练集和测试集
# =========================
X_train, X_test, y_train, y_test = train_test_split(
texts,
labels,
test_size=0.25,
random_state=42,
stratify=labels
)
# =========================
# 5. 定义模型
# =========================
model = Pipeline([
("tfidf", TfidfVectorizer(
tokenizer=chinese_tokenizer,
token_pattern=None
)),
("classifier", LogisticRegression(max_iter=1000))
])
# =========================
# 6. 训练模型
# =========================
model.fit(X_train, y_train)
# =========================
# 7. 测试模型
# =========================
y_pred = model.predict(X_test)
print("\n测试集预测结果：")
for text, true_label, pred_label in zip(X_test, y_test, y_pred):
print(f"文本：{text}")
print(f"真实分类：{true_label}")
print(f"预测分类：{pred_label}")
print("-" * 30)
print("\n分类报告：")
print(classification_report(y_test, y_pred))
# =========================
# 8. 预测新文本
# =========================
new_texts = [
"React hooks 怎么使用",
"MySQL 慢查询怎么优化",
"什么是大模型微调",
"我今天想出去玩"
]
predictions = model.predict(new_texts)
print("\n新文本预测：")
for text, label in zip(new_texts, predictions):
print(f"{text} => {label}")
# =========================
# 9. 保存模型
# =========================
joblib.dump(model, "text_classifier.pkl")
print("\n模型已保存到 text_classifier.pkl")
4. 解释 Pipeline
核心代码是：
model = Pipeline([
("tfidf", TfidfVectorizer(
tokenizer=chinese_tokenizer,
token_pattern=None
)),
("classifier", LogisticRegression(max_iter=1000))
])
这相当于一条流水线：
原始文本
↓
TF-IDF：把文字变成数字特征
↓
LogisticRegression：根据数字特征分类
↓
输出：前端 / 后端 / AI / 其他
TfidfVectorizer 负责提取文本特征。
比如：
React 组件如何传递 props
会变成类似：
React: 0.43
组件: 0.51
传递: 0.38
props: 0.62
LogisticRegression 则学习：
React、Vue、CSS、JavaScript -> 前端
Spring、MySQL、Redis、接口 -> 后端
神经网络、BERT、PyTorch -> AI
天气、电影、吃饭、旅游 -> 其他
5. model.fit 的作用
model.fit(X_train, y_train)
这行就是训练模型。
可以理解成：
model.fit(题目, 正确答案)
在文本分类里：
X_train = 文本
y_train = 正确分类
训练后，模型会从样本中学习哪些词更可能属于哪个类别。
6. 保存和加载模型
训练结束后保存模型：
joblib.dump(model, "text_classifier.pkl")
之后可以单独写一个 predict.py 来预测：
import jieba
import joblib
def chinese_tokenizer(text):
return list(jieba.cut(text))
model = joblib.load("text_classifier.pkl")
while True:
text = input("请输入要分类的文本，输入 q 退出：")
if text.lower() == "q":
break
pred = model.predict([text])[0]
print("预测分类：", pred)
print("-" * 30)
运行：
python predict.py
这样就可以输入一句文本，让模型预测分类。
第二部分：升级到 Transformer 版本
简单模型跑通后，开始进入 Transformer 版本。
目录结构：
project/
├── data.csv
├── train.py
├── predict.py
├── text_classifier.pkl
└── transformer_version/
├── data.csv
├── train_transformer.py
├── predict_transformer.py
└── saved_model/
进入 transformer_version 目录后，单独维护 Transformer 版本代码。
1. 安装依赖
pip install torch transformers pandas scikit-learn
2. Transformer 版本和 TF-IDF 版本的对应关系
TF-IDF 版本：
文本
↓
TfidfVectorizer
↓
LogisticRegression
↓
分类结果
Transformer 版本：
文本
↓
AutoTokenizer
↓
BERT / Transformer
↓
分类头
↓
分类结果
在 Transformer 里，文本不能直接输入模型，需要先通过 tokenizer 转成：
input_ids
attention_mask
labels
3. 什么是 tensor？
在 PyTorch 里，模型不能直接处理普通 Python 列表或字符串。
它处理的是 tensor。
可以简单理解为：
tensor = PyTorch 里的多维数组
例如：
import torch
x = torch.tensor([1, 2, 3])
在深度学习中，文本、图片、标签、模型参数，最后都会变成 tensor。
tensor 的特点是：
可以做高效数学计算
可以放到 GPU 上
可以自动求梯度
4. 什么是 Dataset？
Transformer 训练时，我们需要把原始文本整理成模型能吃的格式。
所以定义一个 TextDataset：
class TextDataset(Dataset):
def __init__(self, texts, labels, tokenizer, max_length):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, index):
text = self.texts[index]
label = self.labels[index]
encoding = self.tokenizer(
text,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt"
)
return {
"input_ids": encoding["input_ids"].squeeze(0),
"attention_mask": encoding["attention_mask"].squeeze(0),
"labels": torch.tensor(label, dtype=torch.long)
}
它的作用是：
原始文本
↓
tokenizer
↓
input_ids / attention_mask / labels
5. tokenizer 这段代码的含义
encoding = self.tokenizer(
text,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt"
)
它的作用是把文本转成数字。
比如：
React 组件如何传递 props
可能会变成：
input_ids = [101, 1234, 5678, 7890, 102, 0, 0, 0, ...]
参数含义：
max_length=64 每条文本最多 64 个 token
padding="max_length" 短文本补齐到 64
truncation=True 长文本超过 64 就截断
return_tensors="pt" 返回 PyTorch tensor
6. input_ids、attention_mask、labels
最终 Dataset 返回：
{
"input_ids": ...,
"attention_mask": ...,
"labels": ...
}
含义如下：
input_ids 文本转成的 token 编号
attention_mask 哪些是真实文本，哪些是 padding
labels 正确分类
比如：
{
"input_ids": tensor([101, 1234, 5678, 7890, 102, 0, 0, 0]),
"attention_mask": tensor([1, 1, 1, 1, 1, 0, 0, 0]),
"labels": tensor(0)
}
其中：
1 表示真实文本
0 表示补齐内容
.squeeze(0) 是为了去掉多余的一层维度。
tokenizer 返回的形状可能是：
(1, 64)
经过：
.squeeze(0)
变成：
(64,)
然后 DataLoader 会自动把多条样本合成 batch：
input_ids.shape = (batch_size, 64)
第三部分：Transformer 训练代码
在 transformer_version/ 下创建 train_transformer.py：
import os
import json
import torch
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from torch.optim import AdamW
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
# =========================
# 1. 基本配置
# =========================
MODEL_NAME = "bert-base-chinese"
DATA_PATH = "data.csv"
SAVE_DIR = "saved_model"
MAX_LENGTH = 64
BATCH_SIZE = 4
EPOCHS = 3
LR = 2e-5
# =========================
# 2. 标签映射
# =========================
label2id = {
"前端": 0,
"后端": 1,
"AI": 2,
"其他": 3
}
id2label = {
0: "前端",
1: "后端",
2: "AI",
3: "其他"
}
# =========================
# 3. 定义 Dataset
# =========================
class TextDataset(Dataset):
def __init__(self, texts, labels, tokenizer, max_length):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, index):
text = self.texts[index]
label = self.labels[index]
encoding = self.tokenizer(
text,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt"
)
return {
"input_ids": encoding["input_ids"].squeeze(0),
"attention_mask": encoding["attention_mask"].squeeze(0),
"labels": torch.tensor(label, dtype=torch.long)
}
# =========================
# 4. 读取数据
# =========================
df = pd.read_csv(DATA_PATH)
texts = df["text"].tolist()
labels = df["label"].tolist()
numeric_labels = [label2id[label] for label in labels]
print("数据总数：", len(texts))
print("类别分布：")
print(df["label"].value_counts())
# =========================
# 5. 划分训练集和测试集
# =========================
train_texts, test_texts, train_labels, test_labels = train_test_split(
texts,
numeric_labels,
test_size=0.25,
random_state=42,
stratify=numeric_labels
)
# =========================
# 6. 加载 Tokenizer 和模型
# =========================
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(
MODEL_NAME,
num_labels=4,
id2label=id2label,
label2id=label2id
)
# =========================
# 7. 准备 DataLoader
# =========================
train_dataset = TextDataset(
train_texts,
train_labels,
tokenizer,
MAX_LENGTH
)
test_dataset = TextDataset(
test_texts,
test_labels,
tokenizer,
MAX_LENGTH
)
train_loader = DataLoader(
train_dataset,
batch_size=BATCH_SIZE,
shuffle=True
)
test_loader = DataLoader(
test_dataset,
batch_size=BATCH_SIZE,
shuffle=False
)
# =========================
# 8. 设置设备
# =========================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print("当前设备：", device)
# =========================
# 9. 训练模型
# =========================
optimizer = AdamW(model.parameters(), lr=LR)
for epoch in range(EPOCHS):
model.train()
total_loss = 0
for batch in train_loader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels_batch = batch["labels"].to(device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels_batch
)
loss = outputs.loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_loader)
print(f"Epoch {epoch + 1}/{EPOCHS}, Loss: {avg_loss:.4f}")
# =========================
# 10. 测试模型
# =========================
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for batch in test_loader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels_batch = batch["labels"].to(device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask
)
logits = outputs.logits
preds = torch.argmax(logits, dim=1)
all_preds.extend(preds.cpu().tolist())
all_labels.extend(labels_batch.cpu().tolist())
print("\n分类报告：")
print(
classification_report(
all_labels,
all_preds,
target_names=["前端", "后端", "AI", "其他"],
zero_division=0
)
)
# =========================
# 11. 保存模型
# =========================
os.makedirs(SAVE_DIR, exist_ok=True)
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)
with open(os.path.join(SAVE_DIR, "label_map.json"), "w", encoding="utf-8") as f:
json.dump(
{
"label2id": label2id,
"id2label": id2label
},
f,
ensure_ascii=False,
indent=2
)
print(f"\n模型已保存到：{SAVE_DIR}")
1. 训练循环解释
核心训练代码：
optimizer = AdamW(model.parameters(), lr=LR)
for epoch in range(EPOCHS):
model.train()
total_loss = 0
for batch in train_loader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels_batch = batch["labels"].to(device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels_batch
)
loss = outputs.loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_loader)
print(f"Epoch {epoch + 1}/{EPOCHS}, Loss: {avg_loss:.4f}")
这段代码相当于手写版的：
model.fit(...)
具体流程：
创建优化器 AdamW
↓
训练 EPOCHS 轮
↓
每轮从 train_loader 里一批一批取数据
↓
把数据放到 CPU / GPU
↓
模型预测
↓
根据 labels 自动计算 loss
↓
清空旧梯度
↓
根据 loss 反向传播，计算梯度
↓
优化器更新参数
↓
统计平均 loss
最核心的是三行：
optimizer.zero_grad()
loss.backward()
optimizer.step()
含义是：
zero_grad：清空旧梯度
backward：根据 loss 计算新梯度
step：根据梯度更新参数
2. loss 和梯度的关系
模型训练本质上是在求函数参数。
可以把模型看成：
预测结果 = f(输入, 参数)
训练要找到一组好的参数，让预测结果尽量接近正确答案。
loss 的作用是衡量当前参数好不好：
loss 大 = 模型错得严重
loss 小 = 模型预测接近正确答案
梯度 的作用是告诉参数应该怎么改：
梯度 = loss 对参数的变化方向
训练过程就是：
随机初始化参数
↓
用当前参数做预测
↓
计算 loss
↓
根据 loss 计算梯度
↓
根据梯度更新参数
↓
重复很多次
↓
得到一组更好的参数
3. Transformer 中的 Multi-Head Attention 在哪里？
之前手写过一个多头注意力模块：
class MultiHeadAttention(nn.Module):
...
它是 Transformer 内部的核心零件。
而这行代码：
model = AutoModelForSequenceClassification.from_pretrained(
MODEL_NAME,
num_labels=4,
id2label=id2label,
label2id=label2id
)
加载的是一个完整的 BERT 分类模型。
它内部大概是：
BERT 文本分类模型
├── Embedding 层
├── Transformer Layer 1
│ ├── MultiHeadAttention
│ ├── FeedForward
│ ├── LayerNorm
│ └── 残差连接
├── Transformer Layer 2
│ ├── MultiHeadAttention
│ ├── FeedForward
│ ├── LayerNorm
│ └── 残差连接
├── ...
├── Transformer Layer 12
│ ├── MultiHeadAttention
│ ├── FeedForward
│ ├── LayerNorm
│ └── 残差连接
└── classifier 分类头
所以手写的 MultiHeadAttention 对应的是：
AutoModelForSequenceClassification 内部 Transformer Layer 的 attention 部分
在训练代码中，它发生在这一行内部：
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels_batch
)
这行表面上很短，但内部经历了：
input_ids
↓
Embedding
↓
多层 Transformer
↓
MultiHeadAttention
↓
FeedForward
↓
classifier
↓
logits / loss
第四部分：运行训练
进入目录：
cd transformer_version
运行：
python train_transformer.py
训练完成后会输出分类报告，例如：
分类报告：
precision recall f1-score support
前端 1.00 0.33 0.50 3
后端 0.40 1.00 0.57 2
AI 1.00 0.67 0.80 3
其他 1.00 1.00 1.00 2
accuracy 0.70 10
macro avg 0.85 0.75 0.72 10
weighted avg 0.88 0.70 0.70 10
这个结果表示：
测试集一共 10 条
预测对了 7 条
accuracy = 0.70
训练完成后，会保存模型：
模型已保存到：saved_model
目录中会出现：
transformer_version/
├── data.csv
├── train_transformer.py
├── predict_transformer.py
└── saved_model/
├── config.json
├── model.safetensors
├── tokenizer.json
├── vocab.txt
└── label_map.json
第五部分：Transformer 预测脚本
训练完成后，写一个单独的预测脚本。
创建 predict_transformer.py：
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# =========================
# 1. 基本配置
# =========================
MODEL_DIR = "saved_model"
MAX_LENGTH = 64
# =========================
# 2. 加载模型和 tokenizer
# =========================
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
print("模型加载完成，当前设备：", device)
# =========================
# 3. 定义预测函数
# =========================
def predict(text):
encoding = tokenizer(
text,
max_length=MAX_LENGTH,
padding="max_length",
truncation=True,
return_tensors="pt"
)
input_ids = encoding["input_ids"].to(device)
attention_mask = encoding["attention_mask"].to(device)
with torch.no_grad():
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask
)
logits = outputs.logits
probs = torch.softmax(logits, dim=1)
pred_id = torch.argmax(probs, dim=1).item()
label = model.config.id2label[pred_id]
return label, probs[0].cpu().tolist()
# =========================
# 4. 命令行交互预测
# =========================
while True:
text = input("\n请输入要分类的文本，输入 q 退出：")
if text.lower() == "q":
break
label, probs = predict(text)
print("预测分类：", label)
print("各类别概率：")
for i, prob in enumerate(probs):
class_name = model.config.id2label[i]
print(f"{class_name}: {prob:.4f}")
运行：
python predict_transformer.py
测试输入：
React hooks 怎么使用
MySQL 慢查询怎么优化
什么是大模型微调
我今天晚上吃什么
输出类似：
预测分类：前端
各类别概率：
前端: 0.7213
后端: 0.1432
AI: 0.0981
其他: 0.0374
第六部分：总结
这次完整跑通了两个版本的文本分类模型。
TF-IDF 版本
优点：
简单
快
容易理解
适合入门
缺点：
主要依赖关键词
语义理解能力弱
数据少时不稳定
Transformer 版本
优点：
能利用预训练语言模型
语义理解能力更强
适合真实文本分类任务
缺点：
代码复杂
训练慢
最好有 GPU
小数据集上也可能不稳定
最终完整流程是：
data.csv
↓
读取文本和标签
↓
tokenizer 转成 input_ids / attention_mask
↓
BERT 文本分类模型
↓
计算 loss
↓
反向传播更新参数
↓
测试模型
↓
保存 saved_model
↓
predict_transformer.py 加载模型并预测
到这里，一个完整的 Transformer 文本分类模型就跑通了。
相关package：
pip install pandas scikit-learn jieba
pip install joblib
Transformer版本
pip install torch transformers scikit-learn
评论
填写昵称与邮箱即可评论，无需登录。