如何在TFDF库的随机森林模型中添加类别权重处理少数类

针对TensorFlow Decision Forests（TFDF）中处理类别不平衡的需求，以下是几种可行的实现方案，替代sklearn中直接设置class_weight的方式：

1. 样本加权法（推荐，灵活复刻sklearn逻辑） #

TFDF支持通过样本权重间接实现类别权重的效果，核心思路是给少数类样本赋予更高的权重，计算逻辑可与sklearn的class_weight='balanced'对齐：

import tensorflow as tf
import tensorflow_decision_forests as tfdf
import pandas as pd

# 加载数据集（示例）
df = pd.read_csv("your_dataset.csv")
label_col = "target"

# 计算类别权重，模拟balanced效果
class_counts = df[label_col].value_counts()
total_samples = len(df)
num_classes = len(class_counts)
class_weights = {cls: total_samples/(class_counts[cls]*num_classes) for cls in class_counts.index}

# 为每个样本添加权重列
df["sample_weight"] = df[label_col].map(class_weights)

# 转换为TFDF兼容的数据集
train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(df, label=label_col)

# 初始化模型并指定权重列
model = tfdf.keras.RandomForestModel(
    task=tfdf.keras.Task.CLASSIFICATION,
    sample_weight_key="sample_weight"
)

# 训练模型
model.fit(train_ds)

2. 利用Yggdrasil原生类别权重配置 #

TFDF底层依赖Yggdrasil Decision Forests，可通过AdvancedArguments直接传递类别权重参数：

model = tfdf.keras.RandomForestModel(
    task=tfdf.keras.Task.CLASSIFICATION,
    advanced_arguments=tfdf.keras.AdvancedArguments(
        yggdrasil_training_args={
            "classification": {
                # 按类别顺序设置权重，例如第0类权重1.0，第1类权重5.0
                "class_weights": [1.0, 5.0]
            }
        }
    )
)

# 训练前可查看模型识别的类别顺序，确保权重对应正确
# print(model.make_inspector().classes())
model.fit(train_ds)

3. 数据重采样配合训练 #

若不想使用权重，可先对数据集进行过采样少数类/欠采样多数类，再输入TFDF训练：

def resample_imbalanced_data(ds, label_col_idx=-1):
    # 分离两类样本（假设目标是二分类，标签为0和1）
    class_0 = ds.filter(lambda x, y: y == 0)
    class_1 = ds.filter(lambda x, y: y == 1)
    
    # 过采样少数类，使其数量与多数类匹配
    repeat_times = class_0.cardinality() // class_1.cardinality()
    class_1_resampled = class_1.repeat(repeat_times)
    
    # 合并并打乱数据
    return tf.data.Dataset.sample_from_datasets(
        [class_0, class_1_resampled], weights=[0.5, 0.5]
    ).shuffle(1000)

# 转换数据集并重采样
train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(df, label=label_col)
resampled_ds = resample_imbalanced_data(train_ds)

# 训练模型
model = tfdf.keras.RandomForestModel(task=tfdf.keras.Task.CLASSIFICATION)
model.fit(resampled_ds)

以上三种方法中，样本加权法最灵活，能完全对齐sklearn的类别权重逻辑；底层参数法更贴近TFDF原生支持，但需注意类别顺序匹配。

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