Scikit-learn等ML工具中ID列处理与非预测列设置咨询

安全阿岳

2026-5-25

嘿，这个问题我做项目时也反复琢磨过，来给你梳理下最实用的解决方案和各个工具的情况：

一、scikit-learn中ID列处理与预测结果映射的最佳方法

最稳妥且常用的方式就是提前保存测试集的ID列，训练/预测阶段单独分离特征列，最后再合并结果，具体步骤如下：

分离训练集和测试集的ID与特征
不管是训练还是测试，先把ID列单独提取出来保存，再用剩余列作为模型的输入特征：

import pandas as pd
from sklearn.ensemble import RandomForestClassifier

# 加载数据
train_df = pd.read_csv("train.csv")
test_df = pd.read_csv("test.csv")

# 保存测试集ID（关键！）
test_ids = test_df["ID"].copy()

# 移除ID列，准备特征矩阵
X_train = train_df.drop(["ID", "target"], axis=1)
y_train = train_df["target"]
X_test = test_df.drop("ID", axis=1)

训练模型并预测
这一步和常规流程一样，用处理好的特征矩阵训练模型，然后对测试集预测：
```
model = RandomForestClassifier()
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
```

合并ID与预测结果
用保存好的test_ids和预测结果y_pred生成新的DataFrame，直接得到对应关系：

result_df = pd.DataFrame({
    "ID": test_ids,
    "prediction": y_pred
})
result_df.to_csv("predictions_with_id.csv", index=False)

如果是做交叉验证（比如KFold），记得每次划分时都要保留对应 fold 的ID，避免后续无法映射：

from sklearn.model_selection import KFold

kf = KFold(n_splits=5)
for fold, (train_idx, val_idx) in enumerate(kf.split(X_train)):
    # 提取当前fold的训练/验证ID
    train_fold_ids = train_df.loc[train_idx, "ID"]
    val_fold_ids = train_df.loc[val_idx, "ID"]
    
    # 训练模型并验证
    X_train_fold = X_train.iloc[train_idx]
    y_train_fold = y_train.iloc[train_idx]
    X_val_fold = X_train.iloc[val_idx]
    y_val_fold = y_train.iloc[val_idx]
    
    model.fit(X_train_fold, y_train_fold)
    val_pred = model.predict(X_val_fold)
    
    # 合并当前fold的验证ID与预测结果
    val_result = pd.DataFrame({
        "ID": val_fold_ids,
        "prediction": val_pred,
        "actual": y_val_fold
    })

二、scikit-learn是否支持设置“非预测列”？

直接说结论：scikit-learn的核心模型（estimator）本身不支持直接标记“非预测列”。它的API设计是基于纯特征矩阵（X）和目标向量（y）的，X必须全部是用于训练的特征列，不能包含ID这种无关列。

不过可以借助ColumnTransformer来更规范地处理列筛选，比如明确指定哪些列是特征，哪些要排除：

from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder

# 假设我们有数值特征列和类别特征列，ID列不在其中
numerical_features = ["age", "income"]
categorical_features = ["gender", "occupation"]

# 定义转换器，只处理指定的特征列，其余列（包括ID）会被丢弃
preprocessor = ColumnTransformer(
    transformers=[
        ("num", StandardScaler(), numerical_features),
        ("cat", OneHotEncoder(), categorical_features)
    ],
    remainder="drop"  # 丢弃未指定的列，比如ID
)

# 预处理后得到纯特征矩阵
X_train_processed = preprocessor.fit_transform(X_train)
X_test_processed = preprocessor.transform(X_test)

本质上还是手动分离，只是用工具类更系统化而已。

三、其他主流机器学习工具的支持情况

TensorFlow/Keras

TensorFlow本身也没有专门的“非预测列”标记功能，但可以通过两种方式灵活处理：

用tf.feature_column定义特征：可以把ID列定义为一个不参与训练的列（比如tf.feature_column.numeric_column("ID")），但在模型输入时只把特征列连接到输出层，ID列仅作为输入的一部分保留，预测时再提取出来合并结果。

用tf.data分离ID与特征：把ID和特征、目标分开加载，比如：

# 从DataFrame创建数据集，ID单独作为一个元素
dataset = tf.data.Dataset.from_tensor_slices(
    (test_df["ID"].values, test_df.drop("ID", axis=1).values)
)
# 预测时遍历数据集，同时获取ID和预测结果
for id_val, features in dataset.batch(32):
    pred = model.predict(features)
    # 这里可以把id_val和pred对应保存

Spark ML

Spark ML在这方面的设计非常贴合业务场景，原生支持保留非预测列！因为它的Pipeline是基于DataFrame操作的：

你只需要用VectorAssembler把所有特征列组装成一个features列，ID等非特征列会一直保留在DataFrame中；
当你用Pipeline训练模型并调用transform()方法时，预测结果会作为新列（比如prediction）添加到原DataFrame里，直接就能得到ID和预测结果的对应关系：

from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml import Pipeline

# 定义特征列，排除ID和target
feature_cols = [col for col in train_df.columns if col not in ["ID", "target"]]

# 组装特征列
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")
rf = RandomForestClassifier(labelCol="target", featuresCol="features")

# 构建Pipeline
pipeline = Pipeline(stages=[assembler, rf])
model = pipeline.fit(train_df)

# 预测后，结果DataFrame包含ID、原特征列、features列、prediction列
predictions = model.transform(test_df)

# 只保留ID和预测结果
result_df = predictions.select("ID", "prediction")

这种方式完全不需要手动保存和合并ID，非常省心。

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

Scikit-learn等ML工具中ID列处理与非预测列设置咨询

TensorFlow/Keras #

Spark ML #

TensorFlow/Keras

Spark ML