泰坦尼克数据集逻辑回归代码中predict_proba与[:,1]含义解析

predict_proba函数的作用 #

• 这是scikit-learn分类模型自带的方法，核心作用是输出每个样本属于各个类别的概率值
• 针对泰坦尼克的二分类任务（预测是否存活：通常0代表死亡，1代表存活），它会返回一个形状为[测试样本数, 2]的二维数组
• 数组里的每一行对应一个测试样本：第一列（索引为0）是该样本属于“死亡”类别的概率，第二列（索引为1）是属于“存活”类别的概率，每行两个概率的和固定为1

[:,1]的含义 #

• 这是NumPy数组的索引语法，作用是提取二维数组中所有行的第2列（索引为1）的数据
• 绘制ROC曲线时，roc_curve函数需要模型对正类（这里是“存活”，对应类别1）的预测概率作为输入，所以我们要把所有样本的存活概率单独提取出来，这就是[:,1]的作用
• 如果你的任务里把“死亡”设为正类，那这里就会用[:,0]，但泰坦尼克任务中一般以“存活”作为关注的正类，所以用[:,1]

举个简单例子，假设predict_proba(X_test)返回的数组是：

[[0.3, 0.7],
 [0.8, 0.2],
 [0.1, 0.9]]

那么[:,1]提取后得到的就是[0.7, 0.2, 0.9]，也就是三个样本的存活概率。

内容的提问来源于stack exchange，提问作者S_Naghiyeva