如何为XGBRegressor添加重要特征属性并随JSON模型保存

问题1：消除getattr触发的B010和B009警告 #

警告原因 #

• B009：函数参数已通过类型注解model: XGBRegressor明确类型，额外的isinstance检查属于冗余代码
• B010：getattr未指定默认值，属性不存在时会抛出AttributeError，同时不符合代码检查规范

优化方案 #

方案1：简化现有函数

移除冗余的isinstance检查，给getattr添加默认值None，直接解决两个警告：

from typing import List, Optional
from xgboost import XGBRegressor

def get_model_features(model: Optional[XGBRegressor]) -> Optional[List[str]]:
    return getattr(model, "selected_features", None) if model is not None else None

方案2：子类化XGBRegressor（更规范）

动态添加属性容易引发类型检查问题，子类化可显式声明属性，让类型检查工具直接识别：

from typing import List, Optional
from xgboost import XGBRegressor

class CustomXGBRegressor(XGBRegressor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.selected_features: Optional[List[str]] = None

# 训练并设置属性
regressor = CustomXGBRegressor()
regressor.fit(X=X_train, y=y_train, eval_set=[(X_train, y_train), (X_validation, y_validation)], verbose=False)

feature_importance = sorted(
    regressor.get_booster().get_score(importance_type="gain").items(), key=lambda x: x[1]
)
regressor.selected_features = [x[0] for x in feature_importance if x[1] > 0]

# 获取属性无需getattr，直接访问
def get_model_features(model: Optional[CustomXGBRegressor]) -> Optional[List[str]]:
    return model.selected_features if model is not None else None

问题2：自定义属性随JSON模型保存（无需pickle） #

XGBoost的Booster对象支持通过set_attr和get_attr存储自定义字符串属性，这些属性会被写入JSON模型文件。只需将特征列表序列化为字符串即可实现存储：

保存模型时存储属性 #

import json
from xgboost import XGBRegressor

# 训练模型并筛选特征
regressor.fit(X=X_train, y=y_train, eval_set=[(X_train, y_train), (X_validation, y_validation)], verbose=False)

feature_importance = sorted(
    regressor.get_booster().get_score(importance_type="gain").items(), key=lambda x: x[1]
)
selected_features = [x[0] for x in feature_importance if x[1] > 0]

# 将特征列表序列化为JSON字符串，存入Booster属性
regressor.get_booster().set_attr(selected_features=json.dumps(selected_features))

# 保存模型
regressor.save_model("model.json")

加载模型时恢复属性 #

import json
from xgboost import XGBRegressor

# 加载模型
regressor = XGBRegressor()
regressor.load_model("model.json")

# 读取并反序列化属性
selected_features_str = regressor.get_booster().get_attr("selected_features", default=None)
selected_features = json.loads(selected_features_str) if selected_features_str else None

# 可选：将属性赋值回模型实例，方便后续直接访问
setattr(regressor, "selected_features", selected_features)

原理说明 #

XGBoost的JSON模型文件会保存Booster的所有自定义属性（通过set_attr设置），这些属性以字符串形式存储。因此需将列表等复杂结构序列化为字符串，加载时再反序列化，全程无需使用pickle。

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