如何修改 train.py，将 confusion matrix 和模型保存到 outputs/？

2026-11-04技术老齐

在 train.py 中创建 outputs/ 目录，用 fig.savefig() 保存混淆矩阵图片，用 joblib.dump() 保存训练好的模型。保存操作应放在模型完成 fit() 并生成测试集预测结果之后。

建议将输出目录设在 train.py 所在文件夹下，这样从不同目录启动脚本时，保存位置也不会改变。

修改步骤

用 Path.mkdir(parents=True, exist_ok=True) 创建输出目录，目录已存在也能继续运行。
根据 y_test 和 y_pred 计算混淆矩阵，绘图并保存为 PNG。
用 joblib.dump() 保存训练好的模型。

如果尚未安装依赖，可在项目的 Python 环境中运行：

python -m pip install scikit-learn matplotlib joblib

完整的 train.py 示例

示例使用 scikit-learn 自带的 Iris 数据集和决策树分类器。如果已有训练代码，可以保留原来的模型和数据处理流程，加入创建目录、绘图和保存模型的代码。

from pathlib import Path

import joblib
import matplotlib

# 只生成图片，适用于没有图形界面的运行环境。
# 必须放在导入 matplotlib.pyplot 之前。
matplotlib.use("Agg")

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier


def main():
    # outputs/ 位于 train.py 所在目录下。
    output_dir = Path(__file__).resolve().parent / "outputs"
    output_dir.mkdir(parents=True, exist_ok=True)

    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data,
        iris.target,
        test_size=0.2,
        random_state=42,
        stratify=iris.target,
    )

    model = DecisionTreeClassifier(random_state=42)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)

    # 固定类别顺序，使矩阵与显示的类别名称一一对应。
    labels = list(range(len(iris.target_names)))
    cm = confusion_matrix(y_test, y_pred, labels=labels)

    fig, ax = plt.subplots(figsize=(6, 5))
    display = ConfusionMatrixDisplay(
        confusion_matrix=cm,
        display_labels=iris.target_names,
    )
    display.plot(ax=ax, cmap="Blues", values_format="d")
    ax.set_title("Iris confusion matrix")
    fig.tight_layout()

    matrix_path = output_dir / "confusion_matrix.png"
    fig.savefig(matrix_path, dpi=150, bbox_inches="tight")
    plt.close(fig)

    model_path = output_dir / "iris_model.joblib"
    joblib.dump(model, model_path)

    print(f"Confusion matrix saved to: {matrix_path}")
    print(f"Model saved to: {model_path}")


if __name__ == "__main__":
    main()

在 train.py 所在目录运行：

python train.py

脚本成功执行后，目录结构如下：

project/
├── train.py
└── outputs/
    ├── confusion_matrix.png
    └── iris_model.joblib

Path(__file__).resolve().parent 指向脚本所在目录。如果脚本位于 project/src/train.py，输出目录就是 project/src/outputs/。如果要保存到项目根目录，需按实际目录结构调整路径，例如使用 .parent.parent / "outputs"。

如何接入已有训练代码

示例中的 model、X_test 和 y_test 分别是训练好的模型、测试集特征和测试集真实标签。接入现有脚本时，将它们替换为项目中的变量名即可。

类别标签也要与数据一致。示例中的 iris.target_names 对应 Iris 数据集的整数标签。如果项目使用字符串标签或自定义编码，需同步修改 labels 和 display_labels，确保顺序一致。

混淆矩阵的行表示真实类别，列表示预测类别。评估时应使用留出的测试集，避免把训练集上的结果当作模型对新数据的表现。

加载模型与注意事项

在 train.py 所在目录启动 Python 后，可以这样加载模型：

from pathlib import Path

import joblib

model = joblib.load(Path("outputs") / "iris_model.joblib")
# predictions = model.predict(X_new)

X_new 的特征数量和排列顺序必须与训练时一致。如果训练流程包含标准化、缺失值填充等预处理，建议将预处理步骤和分类器组成 Pipeline，保存整个 Pipeline，以便预测时使用相同的处理流程。

使用时还需注意：

同名文件会被覆盖。如果需要保留多次训练结果，可在文件名中加入时间或实验编号。
保存图片后调用 plt.close(fig)，释放图形资源。这里只需写入 PNG，无须调用 plt.show()。
只用 joblib.load() 加载可信文件，加载恶意文件可能执行任意代码。
保存和加载模型时应保持依赖环境一致，尤其是 scikit-learn 版本。跨版本加载不保证兼容。