如何编写条件表达式解决XGBoost预测概率映射键值错位问题

问题原因 #

1. XGBoost分类器的predict_proba()返回的概率数组顺序，和训练时标签的唯一值排序结果完全对应，不是固定匹配1-8的数值。第二个模型训练时没有类别4，返回的概率数组长度为6，对应顺序是[1,2,3,5,6,7]，你原来的逻辑按index从1自增赋值，就把第4位概率错误匹配到了不存在的类别4，导致映射错位。
2. 你后续调整的循环逻辑存在bug：当index=4时，判断不在目标类别列表就执行continue，但index +=1写在了continue之后，永远不会被执行，index一直卡在4，循环无法继续。

解决方案 #

优先使用XGBoost自带的classes_属性做映射，该属性存储了模型训练用到的所有类别，顺序和predict_proba返回的概率顺序完全匹配，天然适配8类和6类两种场景，不需要额外做边界判断：

import numpy as np
from collections import OrderedDict
import operator

# 原有映射字典保持不变
my_dict ={1:'a', 2:'b', 3:'c', 4:'d', 5:'e', 6:'f', 7:'g', 8:'f'}

inputData = [1, 2, 3, 4, 5]
inputData2 = np.array([inputData])

# 获取概率数组和对应的类别顺序
proba_list = xgb_model.predict_proba(inputData2, ntree_limit=None, validate_features=False, base_margin=None)[0]
class_list = xgb_model.classes_

result_data = OrderedDict()
# 直接按类别和概率的对应关系赋值
for class_val, proba in zip(class_list, proba_list):
    result_data[my_dict[class_val]] = round(proba, 2)

# 打印逻辑保持不变
print("result_name : ", max(result_data.items(), key=operator.itemgetter(1))[0])
print("result_value : ", max(proba_list))
print(result_data)

如果你要保留原来的索引遍历逻辑，只需要把index +=1移到判断逻辑外部，避免索引卡住即可：

index = 1
result_data = OrderedDict()
proba_list = xgb_model.predict_proba(inputData2,ntree_limit=None, validate_features=False,base_margin=None)[0]
unique_target = y.target.unique().tolist()
for x in proba_list:
    if index in unique_target:
        result_data[func(index)] = round(x,2)
    # 不管是否匹配到类别，索引都自增
    index += 1

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