机器学习实战：用Python手把手教你绘制ROC曲线（附完整代码）

在机器学习模型评估中，我们常常需要更直观的工具来理解模型性能。想象一下这样的场景：你训练了一个二分类模型，准确率看起来不错，但当实际部署时却发现对某些关键样本的识别效果不佳。这时候，仅靠单一指标就像用一把尺子测量整个房间——它能告诉你长度，却无法展现空间的全貌。ROC曲线正是解决这一痛点的利器，它通过动态阈值下的性能变化，为我们提供了模型表现的立体视角。

本文将摒弃复杂的数学推导，直接从代码实战出发，带你用Python实现ROC曲线的完整绘制流程。无论你是刚入门的新手还是希望巩固基础的中级开发者，都能通过本文掌握这一核心技能。我们将重点解决样本不均衡、阈值选择等实际问题，并提供可直接复用的代码模块。

1. 环境准备与数据加载

在开始绘制ROC曲线之前，我们需要准备好Python环境和示例数据集。推荐使用Anaconda创建独立的虚拟环境，避免包版本冲突：

conda create -n roc_demo python=3.8
conda activate roc_demo
pip install numpy pandas matplotlib scikit-learn

对于演示数据，我们使用sklearn自带的乳腺癌数据集。这个数据集非常适合二分类任务，且包含适度的特征维度：

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split

# 加载数据集
data = load_breast_cancer()
X, y = data.data, data.target

# 查看类别分布
print(f"正样本数量: {sum(y==1)}")
print(f"负样本数量: {sum(y==0)}")

# 划分训练测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

注意：在实际项目中遇到样本不均衡时，可以考虑使用过采样（如SMOTE）或欠采样技术，但本文为保持焦点将不做处理。

2. 模型训练与概率预测

ROC曲线的绘制需要模型输出的概率值而非简单标签。我们以逻辑回归为例，展示如何获取概率预测：

from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

# 数据标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 训练模型并获取概率
model = LogisticRegression(max_iter=1000)
model.fit(X_train_scaled, y_train)
y_scores = model.predict_proba(X_test_scaled)[:, 1]  # 取正类概率

关键点说明：

• predict_proba返回的是样本属于各个类别的概率矩阵
• 我们只取第二列（索引1）作为正类的概率分数
• 对于不支持概率输出的模型（如SVM），可以使用decision_function

3. ROC曲线绘制核心逻辑

理解ROC曲线的绘制原理至关重要。我们需要计算不同阈值下的TPR（真正例率）和FPR（假正例率）：

from sklearn.metrics import roc_curve
import matplotlib.pyplot as plt

# 计算ROC曲线关键点
fpr, tpr, thresholds = roc_curve(y_test, y_scores)

# 绘制基础曲线
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, label='Logistic Regression (AUC = %0.2f)' % roc_auc)
plt.plot([0, 1], [0, 1], 'k--')  # 绘制对角线
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.show()

这段代码生成的图形中，曲线越靠近左上角表示模型性能越好。对角线代表随机猜测的性能基线。

4. 阈值选择策略与代码实现

阈值选择直接影响模型在实际应用中的表现。下面我们实现一个可视化阈值选择工具：

def plot_roc_with_thresholds(fpr, tpr, thresholds):
    plt.figure(figsize=(10, 6))
    
    # 绘制ROC曲线
    plt.plot(fpr, tpr, color='darkorange', lw=2)
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    
    # 标记特定阈值点
    threshold_points = [0.2, 0.5, 0.8]
    colors = ['green', 'red', 'purple']
    
    for thresh, color in zip(threshold_points, colors):
        # 找到最接近的阈值索引
        idx = np.argmin(np.abs(thresholds - thresh))
        plt.scatter(fpr[idx], tpr[idx], color=color, s=100,
                   label=f'Threshold={thresh:.2f}\n(FPR={fpr[idx]:.2f}, TPR={tpr[idx]:.2f})')
    
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve with Threshold Markers')
    plt.legend(loc="lower right")
    plt.grid(True)
    plt.show()

plot_roc_with_thresholds(fpr, tpr, thresholds)

实际应用中，阈值选择应考虑业务场景：

• 高精度需求：选择较高阈值（如0.8），减少假阳性
• 高召回需求：选择较低阈值（如0.3），尽可能捕捉所有正例
• 平衡需求：选择最靠近左上角的点，或使用Youden指数（TPR-FPR）最大化的阈值

5. 样本不均衡处理实战

当数据存在严重不均衡时，ROC曲线可能出现过于乐观的假象。我们通过人工创建不均衡数据来演示：

from sklearn.utils import resample

# 创建不均衡数据集（正:负 = 1:20）
X_majority = X[y == 0]
X_minority = X[y == 1]
y_majority = y[y == 0]
y_minority = y[y == 1]

# 下采样多数类
X_majority_downsampled = resample(X_majority, 
                                 replace=False,
                                 n_samples=len(X_minority),
                                 random_state=42)
X_balanced = np.vstack((X_majority_downsampled, X_minority))
y_balanced = np.hstack((np.zeros(len(X_minority)), np.ones(len(X_minority))))

# 重新训练模型
X_train, X_test, y_train, y_test = train_test_split(X_balanced, y_balanced, test_size=0.3)
model.fit(X_train, y_train)
y_scores = model.predict_proba(X_test)[:, 1]
fpr, tpr, _ = roc_curve(y_test, y_scores)

# 绘制对比图
plt.plot(fpr, tpr, label='Balanced Data')
plt.plot(fpr_imbalanced, tpr_imbalanced, label='Original Imbalanced Data')
plt.legend()

处理不均衡数据的常见方法对比：

方法	优点	缺点	适用场景
过采样	保留所有信息	可能过拟合	数据量较小时
欠采样	计算效率高	丢失信息	数据量大时
类别权重	无需修改数据	不改变决策边界	模型支持时
合成采样	平衡数据分布	可能生成噪声	中度不均衡

6. 多模型对比与AUC计算

比较不同模型的ROC曲线能直观展示性能差异。我们扩展前面的代码：

from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.metrics import auc

# 初始化多个模型
models = {
    "Logistic Regression": LogisticRegression(max_iter=1000),
    "Random Forest": RandomForestClassifier(n_estimators=100),
    "SVM": SVC(probability=True, kernel='rbf')
}

plt.figure(figsize=(10, 8))

# 训练并绘制每个模型的ROC曲线
for name, model in models.items():
    model.fit(X_train_scaled, y_train)
    y_scores = model.predict_proba(X_test_scaled)[:, 1]
    fpr, tpr, _ = roc_curve(y_test, y_scores)
    roc_auc = auc(fpr, tpr)
    plt.plot(fpr, tpr, label=f'{name} (AUC = {roc_auc:.2f})')

# 完善图形
plt.plot([0, 1], [0, 1], 'k--')
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve Comparison')
plt.legend(loc="lower right")
plt.show()

AUC（曲线下面积）是量化模型性能的重要指标：

• 0.9-1.0：优秀
• 0.8-0.9：良好
• 0.7-0.8：一般
• 0.6-0.7：较差
• 0.5-0.6：无效模型

7. 高级应用与性能优化

对于需要更高性能的场景，我们可以从以下几个方向优化：

并行计算加速：

from joblib import Parallel, delayed

def compute_roc_chunk(chunk, model, X, y):
    # 分块计算ROC
    return roc_curve(y[chunk], model.predict_proba(X[chunk])[:, 1])

# 并行计算
results = Parallel(n_jobs=4)(delayed(compute_roc_chunk)(chunk, model, X_test, y_test) 
                            for chunk in np.array_split(range(len(X_test)), 4))

自定义阈值采样： 当数据量极大时，可以不等距采样阈值来减少计算量：

def custom_threshold_sampling(y_true, y_scores, n_thresholds=100):
    # 在概率密度高的区域密集采样
    hist, bin_edges = np.histogram(y_scores, bins=n_thresholds)
    thresholds = bin_edges[:-1]  # 使用直方图边界作为阈值
    return roc_curve(y_true, y_scores, drop_intermediate=False)

实时ROC监控： 对于在线学习系统，可以实现增量式ROC计算：

class IncrementalROCCalculator:
    def __init__(self):
        self.all_scores = []
        self.all_labels = []
    
    def update(self, y_true_batch, y_scores_batch):
        self.all_scores.extend(y_scores_batch)
        self.all_labels.extend(y_true_batch)
    
    def get_current_roc(self):
        return roc_curve(self.all_labels, self.all_scores)

在实际项目中，我发现ROC曲线与PR曲线（精确率-召回率曲线）结合使用效果最佳。特别是当正样本比例低于10%时，PR曲线往往能更敏感地反映模型性能变化。