如何在Python中pickle/保存整套缺失值插补处理流程

方法1：封装为自定义类后用pickle直接序列化 #

你需要先把零散的处理逻辑、训练得到的中间变量（比如训练好的MissForest实例、每个列的唯一值列表、要保留的列名等）全部封装到一个统一的类里，把训练逻辑放在fit方法，新数据处理逻辑放在transform方法，之后直接序列化这个类的实例即可，示例代码如下：

import pickle
import pandas as pd
from missforest import MissForest

class CustomImputer:
    def __init__(self, drop_cols=['id'], unique_threshold=2000):
        self.drop_cols = drop_cols
        self.unique_threshold = unique_threshold
        self.miss_forest = MissForest()
        self.col_unique_values = {} # 存储训练集每个列的唯一值

    @staticmethod
    def _find_closest_value(target, unique_values):
        chosen = unique_values[0]
        min_l2 = (target - chosen) ** 2
        for value in unique_values:
            current_l2 = (target - value) ** 2
            if current_l2 < min_l2:
                chosen = value
                min_l2 = current_l2
        return chosen

    def fit(self, data):
        # 训练MissForest
        self.feat_cols = [col for col in data.columns if col not in self.drop_cols]
        imputed_data = self.miss_forest.fit_transform(data)
        imputed_data = pd.DataFrame(imputed_data, columns=data.columns)
        # 存储每个列训练集非空唯一值
        for col in self.feat_cols:
            self.col_unique_values[col] = data[col].dropna().unique().tolist()
        return self

    def transform(self, data):
        # 新数据插补流程
        imputed_data = self.miss_forest.transform(data)
        imputed_data = pd.DataFrame(imputed_data, columns=data.columns)
        # 匹配最近唯一值
        for col in self.feat_cols:
            unique_vals = self.col_unique_values[col]
            if len(unique_vals) < self.unique_threshold:
                imputed_data[col] = imputed_data[col].apply(lambda x: self._find_closest_value(x, unique_vals))
        # 删除指定列
        imputed_data = imputed_data.drop(self.drop_cols, axis=1)
        return imputed_data

# 训练+保存
imputer = CustomImputer()
imputer.fit(train_data)
with open("custom_imputer.pkl", "wb") as f:
    pickle.dump(imputer, f)

# 后续加载使用
with open("custom_imputer.pkl", "rb") as f:
    loaded_imputer = pickle.load(f)
new_imputed_data = loaded_imputer.transform(new_data)

注：你原代码里存在循环顺序bug（先遍历columns_name_lst再给它赋值），上述封装代码已经做了修正。

方法2：用joblib序列化（更适合机器学习相关对象） #

joblib是专门针对Python大数据对象优化的序列化工具，对sklearn风格的模型、大数组的存储效率比pickle高很多，用法和pickle几乎一致：

import joblib

# 保存（imputer为上述CustomImputer的训练后实例）
joblib.dump(imputer, "custom_imputer.joblib")

# 加载
loaded_imputer = joblib.load("custom_imputer.joblib")

方法3：拆分保存状态+逻辑（兼容性最好） #

如果担心pickle/joblib的版本兼容问题（比如不同Python版本、不同依赖库版本加载失败），可以把训练得到的所有状态参数单独保存，处理逻辑单独写成固定的函数：

• 要保存的状态包括：训练好的MissForest实例、每个列的唯一值列表、要删除的列名单、唯一值阈值参数
• 可以用joblib存模型参数，用JSON存列名、唯一值列表这种结构化数据

import json

# 保存状态
joblib.dump(imputer.miss_forest, "miss_forest.joblib")
state = {
    "drop_cols": imputer.drop_cols,
    "unique_threshold": imputer.unique_threshold,
    "col_unique_values": imputer.col_unique_values,
    "feat_cols": imputer.feat_cols
}
with open("imputer_state.json", "w") as f:
    json.dump(state, f)

# 加载的时候先读状态，再调用单独写的插补函数即可

注意事项 #

• pickle和joblib都存在安全风险，不要加载来源不明的序列化文件
• 保存和加载环境的依赖库版本（MissForest、pandas、scikit-learn、Python版本）尽量保持一致，避免出现兼容报错

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