用sklearn运行分类模型，选择AUC最高的模型保存模型权重并绘制AUCROC曲线（以逻辑回归、随机森林、梯度提升、MLP为例）

诸神缄默不语2025-01-20 12:39

文章目录

[1. 导入包](#1. 导入包)
[2. 初始化分类模型](#2. 初始化分类模型)
[3. 训练、测试模型，绘图，保存指标](#3. 训练、测试模型，绘图，保存指标)

1. 导入包

python 复制代码

from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import roc_auc_score,accuracy_score,roc_curve,auc
import joblib
import matplotlib.pyplot as plt

2. 初始化分类模型

python 复制代码

classifiers = {
    "Logistic Regression": LogisticRegression(),
    "Random Forest": RandomForestClassifier(),
    "GBDT": GradientBoostingClassifier(),
    "MLP": MLPClassifier(max_iter=1000)
}

3. 训练、测试模型，绘图，保存指标

在这里省略了数据处理部分，总之X/Y都是np.ndarray对象。f反正你创建一个可写的文件流就行，如果连这个都不会的话参考我写的这篇博文：Python3对象序列化，即处理JSON、XML和文件（持续更新ing...）。
f.close()没写，根据你的需要如果想加就加。

这个逻辑是每次得到AUC最高的模型就画图，其实感觉把模型权重储存下来然后再joblib.load()再画图会更合适......

如果想对每个模型画ROC曲线叠在一张图上的话，在最前面新建画布（plt.figure()），每个模型运行完后都运行一次plt.plot()，不close()就行。

python 复制代码

max_auc = 0
max_acc = 0
best_classifier = ""
# 训练模型
for lr_name, lr in classifiers.items():
    lr.fit(X_train, y_train)

    # 预测
    y_pred = lr.predict(X_test)
    y_pred_proba = lr.predict_proba(X_test)[:, 1]

    # 评估
    auc_score = roc_auc_score(y_test, y_pred_proba)
    acc = accuracy_score(y_test, y_pred)

    if auc_score > max_auc:
        max_auc = auc_score
        max_acc = acc
        best_classifier = lr_name
        joblib.dump(lr, f"model.pkl")

        fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
        roc_auc = auc(fpr, tpr)
        plt.figure()
        plt.plot(
            fpr,
            tpr,
            color="darkorange",
            lw=2,
            label=f"ROC curve (AUC = {roc_auc:.2f})",
        )
        plt.plot(
            [0, 1], [0, 1], color="navy", lw=2, linestyle="--"
        )  # 随机猜测基线
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel("False Positive Rate")
        plt.ylabel("True Positive Rate")
        plt.title("Receiver Operating Characteristic")
        plt.legend(loc="lower right")
        plt.grid()
        plt.savefig("roc.png")
        plt.close()

    f.write(
        f"{lr_name} AUC: {auc_score:.4f}, ACC: {acc:.4f}"
        + "\n"
    )
    f.flush()

f.write(
    f"best_classifier: {best_classifier} AUC: {max_auc:.4f}, ACC: {max_acc:.4f}"
    + "\n"
)
f.flush()