XGBoost.save_model保存模型丢失feature_names，如何恢复？

XGBoost的Booster.save_model()默认仅存储模型的结构与权重参数，不会保存feature_names这类元数据，因此加载后会出现feature_names为None的情况。以下两种方法可以解决这个问题：

方法一：保存完整的XGBClassifier对象 #

直接用pickle或joblib保存整个分类器实例，这样会保留所有训练相关的元数据（包括feature_names）：

import pickle
import xgboost as xgb
import pandas as pd

# 数据准备
data = {'col_1': [3, 2, 1, 0], 'col_2': [1, 2, 3, 4], 'label': [1, 1, 0, 0]}
df = pd.DataFrame.from_dict(data)
y_train = df.pop('label')
x_train = df

# 训练模型
classifier = xgb.XGBClassifier(max_depth=3, learning_rate=0.1, n_estimators=3)
classifier.fit(x_train, y_train)
print(f'训练后特征名: {classifier.get_booster().feature_names}')

# 保存完整模型
filename = 'test_model.pkl'
with open(filename, 'wb') as f:
    pickle.dump(classifier, f)

# 加载模型
with open(filename, 'rb') as f:
    loaded_classifier = pickle.load(f)
print(f'加载后特征名: {loaded_classifier.get_booster().feature_names}')

方法二：手动存储并恢复feature_names #

如果必须使用Booster.save_model()，可以在保存前将feature_names存入模型的自定义属性，加载后再读取恢复：

import xgboost as xgb
import pandas as pd

# 数据准备
data = {'col_1': [3, 2, 1, 0], 'col_2': [1, 2, 3, 4], 'label': [1, 1, 0, 0]}
df = pd.DataFrame.from_dict(data)
y_train = df.pop('label')
x_train = df

# 训练模型
classifier = xgb.XGBClassifier(max_depth=3, learning_rate=0.1, n_estimators=3)
classifier.fit(x_train, y_train)
booster = classifier.get_booster()
print(f'训练后特征名: {booster.feature_names}')

# 将feature_names存入模型属性
booster.set_attr(feature_names=','.join(booster.feature_names))
filename = 'test_model.bst'
booster.save_model(filename)

# 加载模型并恢复feature_names
bst_model = xgb.Booster()
bst_model.load_model(filename)
# 读取属性并拆分还原
feature_names = bst_model.attr('feature_names').split(',')
bst_model.feature_names = feature_names
print(f'加载后特征名: {bst_model.feature_names}')

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