诸神缄默不语-个人CSDN博文目录

1. 导入包

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. 初始化分类模型

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()就行。

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_scoremax_acc = accbest_classifier = lr_namejoblib.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()