使用Hugging Face BERT做多标签分类应用自定义损失时出现TypeError

问题原因 #

报错核心来自三处不匹配：

1. 预训练模型加载逻辑错误，模型输出维度和预期不符
2. 多标签场景下评估指标的选择不匹配，默认准确率无法适配多标签输入格式
3. 数据集转换时标签未转为标准张量，存在ragged类型/object类型问题

修复步骤 #

1. 修正预训练模型加载与特征提取逻辑 #

你当前使用TFAutoModelForSequenceClassification（自带分类头）作为encoder，但实际需要的是纯BERT backbone输出的文本特征，且需要明确取 <[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]> token的输出作为分类特征：

with strategy.scope():
    # 改用不带分类头的基础预训练模型
    encoder = TFAutoModel.from_pretrained(checkpoint)
    model = create_model(encoder)

对应修改create_model函数中特征提取部分：

def create_model(encoder, nb_classes=3, lr=1e-5):
    # 输入定义保持不变
    input_ids = tf.keras.Input(shape=(512,), ragged=False,
                               dtype=tf.int32, name='input_ids')
    input_attention_mask = tf.keras.Input(shape=(512,), ragged=False,
                                          dtype=tf.int32, name='attention_mask')
    # 取BERT最后一层所有token输出后，单独提取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]> token的特征
    output = encoder({'input_ids': input_ids, 
                      'attention_mask': input_attention_mask})[0]
    cls_output = output[:, 0, :]
    Y = tf.keras.layers.BatchNormalization()(cls_output)
    Y = tf.keras.layers.Dense(nb_classes, activation='sigmoid')(Y)

2. 修正损失函数与评估指标 #

不需要手动实现交叉熵损失，TensorFlow内置的BinaryCrossentropy完全适配多标签场景，且做了数值稳定性优化。同时多标签场景不能直接用默认的acc（默认是多分类准确率，需要做argmax），需要指定为二分类准确率：

# 接上面create_model函数的逻辑
    model = tf.keras.Model(inputs=[input_ids, input_attention_mask], 
                           outputs=[Y])
    optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
    # 直接用内置损失，不需要自定义实现
    model.compile(optimizer=optimizer, 
                  loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
                  metrics=[tf.keras.metrics.BinaryAccuracy(name='acc')])
    model.summary()
    return model

3. 修正数据集转换逻辑 #

原代码中标签直接从Hugging Face Dataset中取，是列表格式，转为tf张量时没有显式规整为二维float32张量，导致出现ragged张量的类型错误。修改tf_text_data_prep函数的数据集构造部分：

def tf_text_data_prep(df):
    hugging_ds = Dataset.from_pandas(df)
    tokenized_ds = hugging_ds.map(
                      tokenize_function,
                      batched=True,
                      num_proc=strategy.num_replicas_in_sync,
                      remove_columns=["Text", '__index_level_0__'],
                      load_from_cache_file=True 
                      )
    
    tf_dataset = tokenized_ds.with_format("tensorflow")
    features = {x: tf_dataset[x].to_tensor() for x in tokenizer.model_input_names}
    # 显式把标签列表堆叠为二维张量，转为float32类型
    labels = tf.cast(tf.stack(tf_dataset["label"].to_list()), tf.float32)
    tf_data = tf.data.Dataset.from_tensor_slices((features, labels))
    return tf_data

验证 #

修改完成后，你可以先打印一条数据的标签和模型单步预测的输出，确认两者的维度都是(batch_size, 3)，且类型都是float32，即可正常启动训练。之前的nested_row_splits报错也会因为ragged张量的消失而解决。

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