TensorFlow多head Estimator中二分类头logit_dimension=1的问题

我之前也遇到过类似的TensorFlow 1.x多head Estimator的坑，结合你的问题描述，核心问题大概率出在标签传递格式和二分类头的参数匹配上，咱们一步步拆解：

1. 先明确二分类头的核心要求（你之前的理解是对的，但要注意细节） #

在TensorFlow 1.4中，binary_classification_head确实要求logits的最后一维是1，而不是2。这是因为它底层用的是sigmoid_cross_entropy_with_logits损失，输出是单个概率值（对应正类的概率），而非softmax的多类概率分布。所以你把logits_0改成维度1是正确的操作。

2. 关键问题：Multi-Head场景下的标签必须是字典格式 #

这是你训练结果无效的核心原因！当使用multi_head时，labels参数必须是一个字典，键对应每个head的name，值对应该head的标签张量。如果你的输入函数返回的是单一标签张量，二分类头根本拿不到自己对应的标签，相当于模型一直在用随机标签训练，自然结果和baseline差不多。

举个例子，你的输入函数应该返回这样的格式：

def input_fn():
    # 假设你的数据存在DataFrame中
    df = load_your_data()
    features = {"CAT_XXX": df["CAT_XXX"].values}
    # 重点：labels是字典，键和head的name完全匹配
    labels = {
        "target_3": df["target_3"].values,  # 0/1的整数或bool张量
        "target_2": df["target_2"].values   # 0/1/2的整数张量
    }
    return tf.data.Dataset.from_tensor_slices((features, labels)).batch(32).repeat()

3. 修正Model_fn中的细节匹配 #

确保logits字典的键、head的name、labels字典的键三者完全一致，否则TensorFlow无法正确映射：

def model_fn_multihead(features, labels, mode, params):
    # 定义head时，name要和后续的键严格对应
    head_target3 = tf.contrib.estimator.binary_classification_head(name="target_3")
    head_target2 = tf.contrib.estimator.multi_class_head(n_classes=3, name="target_2")
    
    # 创建多head
    head = tf.contrib.estimator.multi_head([head_target3, head_target2])
    
    # 构建共享网络
    net = tf.feature_column.input_layer(features, params['feature_columns'])
    for idx, units in enumerate(params['hidden_units']):
        net = tf.layers.dense(net, units=units, activation=tf.nn.relu, name=f'fully_connected_{idx}')
    
    # 二分类logits维度为1，多分类为3
    logits_target3 = tf.layers.dense(net, 1, activation=None, name='logits_target3')
    logits_target2 = tf.layers.dense(net, 3, activation=None, name='logits_target2')
    
    # logits字典的键必须和head的name完全一致
    logits = {
        "target_3": logits_target3,
        "target_2": logits_target2
    }
    
    def _train_op_fn(loss):
        # 可以尝试调小学习率，Adagrad在简单任务上0.01可能偏大
        return tf.train.AdagradOptimizer(learning_rate=0.001).minimize(
            loss, global_step=tf.train.get_global_step())
    
    return head.create_estimator_spec(
        features=features,
        labels=labels,
        mode=mode,
        logits=logits,
        train_op_fn=_train_op_fn)

4. 额外调试建议 #

如果还是有问题，可以加一些调试步骤确认：

• 在model_fn中打印labels["target_3"]的形状和前10个值，确认标签正确传递；
• 查看logits_target3的输出均值，训练过程中如果均值一直接近0，说明模型没有更新，大概率是标签没传对；
• 单独用binary_classification_head做一个单任务Estimator，用相同的特征和标签训练，确认独立任务能正常收敛，再对比多head场景的差异。

按照这个思路调整后，你的二分类任务应该能和多分类任务一样正常收敛，毕竟你的测试用例中特征和二分类标签是完全线性相关的，模型应该能轻松达到接近100%的准确率。

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