如何修改TensorFlow单输入BERT文本分类模型适配双文本列输入并解决报错

报错根因 #

你之前的拼接逻辑有问题：两个shape=()的标量字符串输入沿axis=-1拼接后，输出形状会变成(batch_size, 2)，TF Hub的BERT预处理层会把这个二维字符串张量的每一个元素都当成独立的单文本处理，原本32的batch就会被扩展成64，最终输出的logits尺寸是[64,6]，和输入标签的[32,6]无法匹配，就出现了你遇到的维度不匹配错误。

正确改造方案 #

方案1：BERT原生句对输入（推荐） #

BERT原生支持文本对输入，只需要把两个文本用特殊分隔符<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>拼接为单个字符串即可，不需要修改模型结构，适配绝大多数句对分类场景（比如语义匹配、问答匹配）。

步骤1：修改数据输入

训练前先把两个文本列拼接：

# 假设你的两个文本列是text_a和text_b
X_train_merged = X_train['text_a'] + '[SEP]' + X_train['text_b']
X_valid_merged = X_valid['text_a'] + '[SEP]' + X_valid['text_b']

步骤2：训练时直接用拼接后的数据

用原来的单输入模型训练即可：

history = classifier_model.fit(
    x=X_train_merged.values,
    y=y_train_c,
    validation_data=(X_valid_merged.values, y_valid_c),
    epochs=15
)

方案2：双路BERT分别编码后拼接 #

如果你的两个文本需要独立编码再融合，用以下方式修改模型：

import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras.optimizers import Adam

tfhub_handle_encoder = "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/1"
tfhub_handle_preprocess = "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"

def build_classifier_model():
    # 定义两个输入层
    input_a = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text_a')
    input_b = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text_b')
    
    # 共享预处理和BERT编码器
    preprocessing_layer = hub.KerasLayer(tfhub_handle_preprocess, name='preprocessing')
    encoder = hub.KerasLayer(tfhub_handle_encoder, trainable=True, name='BERT_encoder')
    
    # 分别处理两个文本
    encoder_input_a = preprocessing_layer(input_a)
    encoder_output_a = encoder(encoder_input_a)['pooled_output']
    
    encoder_input_b = preprocessing_layer(input_b)
    encoder_output_b = encoder(encoder_input_b)['pooled_output']
    
    # 拼接两个编码结果
    merged = tf.keras.layers.concatenate([encoder_output_a, encoder_output_b], axis=-1)
    net = tf.keras.layers.Dropout(0.1)(merged)
    net = tf.keras.layers.Dense(6, activation='softmax', name='classifier')(net)
    
    # 构建双输入模型
    model = tf.keras.Model([input_a, input_b], net)
    
    loss = tf.keras.losses.CategoricalCrossentropy(from_logits=False)
    metric = tf.metrics.CategoricalAccuracy('accuracy')
    optimizer = Adam(learning_rate=5e-05, epsilon=1e-08, decay=0.01, clipnorm=1.0)
    model.compile(optimizer=optimizer, loss=loss, metrics=metric)
    model.summary()
    return model

训练时传入两个输入：

history = classifier_model.fit(
    x=[X_train['text_a'].values, X_train['text_b'].values],
    y=y_train_c,
    validation_data=([X_valid['text_a'].values, X_valid['text_b'].values], y_valid_c),
    epochs=15
)

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