模型训练出来以后,如果每次预测都要重新训练一遍,就还没有进入真实项目状态。
真实项目更常见的流程是:训练一次,保存下来;下次启动程序时加载模型,直接对新文本做预测。
要保存什么
文本分类项目里,最容易犯的错误是只保存分类器:
joblib.dump(model,"model.joblib")这通常不够。因为模型依赖前面的 TF-IDF 词表和 IDF 权重。
更稳的做法是保存完整 Pipeline:
TfidfVectorizer + LogisticRegression = text_classifier_pipeline.joblib这样加载后,模型知道应该用哪套词表、哪套参数、哪种特征空间。
保存完整 Pipeline
假设上一课已经训练好了pipe:
importjoblib joblib.dump(pipe,"text_classifier_pipeline.joblib")保存成功后,会得到一个模型文件。它里面包含:
- TF-IDF 的词表;
- IDF 权重;
- 模型参数;
- Pipeline 步骤顺序。
这比把向量化器和模型分开保存更不容易出错。
加载模型并预测
加载时很简单:
loaded_pipe=joblib.load("text_classifier_pipeline.joblib")然后对新文本预测:
new_texts=[cut_text("这篇文章讲得很清楚"),cut_text("软件一直闪退,完全没法用"),]pred=loaded_pipe.predict(new_texts)print(pred)这里仍然要注意:如果训练阶段输入的是分词后的文本,预测阶段也要先分词。训练和预测的预处理必须一致,否则模型看到的特征空间会变。
封装成可复用函数
可以把加载和预测写成函数:
importjoblib MODEL_PATH="text_classifier_pipeline.joblib"defload_model():returnjoblib.load(MODEL_PATH)defpredict_texts(model,texts):cut_texts=[cut_text(text)fortextintexts]returnmodel.predict(cut_texts)model=load_model()labels=predict_texts(model,["这个功能很好用"])print(labels)这样以后无论是命令行、网页接口,还是 Streamlit 页面,都可以复用这套逻辑。
如果想输出置信度
很多分类器支持predict_proba():
proba=loaded_pipe.predict_proba(new_texts)print(proba)它会输出每个类别的概率估计。
不过要记住:概率不一定等于真实可信度。尤其是数据少、类别不平衡、模型没校准时,概率只能作为参考。
保存标签和元数据
除了模型文件,建议额外保存一份元数据:
importjson meta={"model_name":"tfidf_logistic_regression","version":"v1.0","labels":list(loaded_pipe.classes_),"created_at":"2026-07-01","text_format":"jieba_cut_with_space",}withopen("model_meta.json","w",encoding="utf-8")asf:json.dump(meta,f,ensure_ascii=False,indent=2)元数据能帮你在几个月后快速知道:这个模型是谁训练的、用的什么输入格式、支持哪些标签。
深度学习模型的保存:PyTorch
上面的 joblib 思路适用于 scikit-learn Pipeline。换到 PyTorch,保存方式有两种,区别在于保存的是“整个对象”还是“只是参数”。
importtorch model=MyClassifier()# 方式 1:保存整个模型,简单但耦合类定义torch.save(model,"full_model.pt")loaded=torch.load("full_model.pt")# 方式 2:只保存参数,推荐torch.save(model.state_dict(),"model_weights.pt")# 加载时先重建模型结构,再灌入参数loaded_model=MyClassifier()loaded_model.load_state_dict(torch.load("model_weights.pt"))loaded_model.eval()# 切换到评估模式,关掉 dropout 等| 方式 | 优点 | 缺点 |
|---|---|---|
torch.save(model) | 一行搞定 | 耦合模型类定义,类名或结构一变就加载不了 |
torch.save(state_dict) | 灵活,只有参数 | 需要重建模型结构再load_state_dict |
推荐用state_dict,因为它只存参数,不绑定具体类定义,迁移和重构时更省事。
训练中断也能续上:检查点
训练耗时久的模型,每隔几个 epoch 存一次检查点,服务器挂了也能从最近的状态恢复,而不是从头再来。
checkpoint={"epoch":epoch,"model_state_dict":model.state_dict(),"optimizer_state_dict":optimizer.state_dict(),"loss":loss.item(),}torch.save(checkpoint,f"checkpoint_epoch_{epoch}.pt")恢复时把这几项分别加载回模型和优化器即可。优化器状态也要存,因为 Adam 等优化器内部还维护着动量一类的中间量。
常见踩坑
只保存模型,没保存向量化器
预测时会丢失词表,导致新文本无法转成和训练时一致的特征。
训练和预测分词规则不同
训练时用 jieba,预测时直接传原句,模型效果会明显不稳定。
路径写死
本地能跑,部署到服务器后找不到文件。建议统一使用配置项或相对项目根目录的路径。
依赖版本变化
joblib保存的是 Python 对象。跨环境加载时,最好记录 scikit-learn、jieba、Python 的版本。
importsklearnimportjiebaimportsysprint(sys.version)print(sklearn.__version__)一个最小目录结构
实际项目可以这样组织:
text-classifier/ train.py predict.py models/ text_classifier_pipeline.joblib model_meta.jsontrain.py负责训练并保存模型,predict.py负责加载模型并预测。这样代码职责会更清楚。
小结
文本分类模型要能复用,保存的重点不是“分类器对象”,而是“完整预测链路”。
更稳的顺序是:
训练 Pipeline -> 评估效果 -> 保存完整 Pipeline -> 保存标签和元数据 -> 加载模型 -> 对新文本预测scikit-learn 用joblib存整个 Pipeline,PyTorch 推荐存state_dict,长时间训练再加检查点机制。到这里,中文文本分类项目已经形成闭环。下一阶段我们回到 PyTorch,理解模型训练背后的自动微分机制。
在线阅读
点击这里阅读博客原文
原文地址:https://bestsdz.xyz/posts/save-load-text-classification-model/