在Google Colab TPU上微调expertBERT时调用model.fit()遇InternalError

我在Google Colab中基于Sentiment140 TFDS数据集，微调TensorFlow Hub上的expertBERT模型，调用model.fit()时触发以下错误：

InternalError: RET_CHECK failure (third_party/tensorflow/core/tpu/graph_rewrite/distributed_tpu_rewrite_pass.cc:1992) arg_shape.handle_type != DT_INVALID  input edge: [id=3880 model_4_keras_layer_11_262798:0 -> cluster_train_function:652]

代码在CPU上运行正常，也查过Stack Overflow上的类似问题，但没找到可行的解决方案。相关代码片段如下：

def create_model():
    bert_model = hub.KerasLayer(BERT_URL, trainable=True)
    preprocessor_model = hub.KerasLayer(PREPROCESSOR_URL)

    input_text = tf.keras.layers.Input(shape=(), dtype=tf.string, name="input_layer")
    bert_inputs = preprocessor_model(input_text)
    bert_outputs = bert_model(bert_inputs)
    drop = tf.keras.layers.Dropout(0.3)(bert_outputs.get("pooled_output"))
    output_layer = tf.keras.layers.Dense(1, activation="sigmoid", name="output_layer")(drop)

    model = tf.keras.Model(inputs=input_text, outputs=output_layer)
    return model


def optimize_ds(dataset, batch_size):
    dataset = tf.data.Dataset.from_tensor_slices(dataset["train"]).take(60000) #JUST A TINY SAMPLE
    dataset_size = dataset.cardinality().numpy()
    dataset = dataset.batch(batch_size)
    dataset = dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
    return dataset, dataset_size

with strategy.scope():
    metrics = tf.metrics.BinaryAccuracy()
    loss_function = tf.keras.losses.BinaryCrossentropy(from_logits=False)

    train_dataset, train_data_size = optimize_ds(in_memory_ds, BATCH_SIZE)
    print(train_data_size)
    steps_per_epoch = train_data_size    
    num_train_steps = steps_per_epoch * EPOCHS
    num_warmup_steps = num_train_steps
    model = create_model()
    print(train_dataset.take(1))

    optimizer = optimization.create_optimizer(
        init_lr=INIT_LR,
        num_train_steps=num_train_steps,
        num_warmup_steps=num_warmup_steps,
        optimizer_type="adamw"
    )

    model.compile(optimizer=optimizer, loss=loss_function, metrics=[metrics])

    early_stop = EarlyStopping(monitor="val_loss", patience=PATIENCE, restore_best_weights=True)

    checkpoint_path = "/content/cp.ckpt"
    checkpoint_dir = os.path.dirname(checkpoint_path)
    options = tf.train.CheckpointOptions(experimental_io_device="/job:localhost")
    cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                        save_weights_only=True,
                                                        save_best_only=True,
                                                        verbose=1, options=options)
    latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir)
    if latest_checkpoint is not None:
        print(latest_checkpoint)
        model.load_weights(latest_checkpoint)

    history = model.fit(train_dataset, epochs=EPOCHS, batch_size=BATCH_SIZE,
                        callbacks=[early_stop, cp_callback])

我已经添加了所有TPU运行所需的样板代码。

可行的修复方案 #

1. 调整数据集处理逻辑

TPU对数据集的格式和处理方式更敏感，避免不必要的内存转换：

• 移除tf.data.Dataset.from_tensor_slices，直接使用TFDS原始数据集进行处理，同时确保提取正确的文本和标签字段：

  def optimize_ds(dataset, batch_size):
      dataset = dataset["train"].take(60000)
      # 映射提取Sentiment140的文本和标签字段（根据数据集结构调整键名）
      def map_func(x):
          return x["text"], x["sentiment"]
      dataset = dataset.map(map_func, num_parallel_calls=tf.data.AUTOTUNE)
      dataset = dataset.batch(batch_size)
      # TPU上建议取消cache，或使用分布式兼容的缓存方式
      dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
      dataset_size = dataset.cardinality().numpy()
      return dataset, dataset_size
  

2. 确保所有模型组件在TPU策略范围内创建

将create_model的定义和调用都放在strategy.scope()内部，避免层初始化脱离TPU上下文：

with strategy.scope():
    metrics = tf.metrics.BinaryAccuracy()
    loss_function = tf.keras.losses.BinaryCrossentropy(from_logits=False)

    # 把模型创建移到策略范围内
    def create_model():
        bert_model = hub.KerasLayer(BERT_URL, trainable=True)
        preprocessor_model = hub.KerasLayer(PREPROCESSOR_URL)
        input_text = tf.keras.layers.Input(shape=(), dtype=tf.string, name="input_layer")
        bert_inputs = preprocessor_model(input_text)
        bert_outputs = bert_model(bert_inputs)
        drop = tf.keras.layers.Dropout(0.3)(bert_outputs.get("pooled_output"))
        output_layer = tf.keras.layers.Dense(1, activation="sigmoid", name="output_layer")(drop)
        return tf.keras.Model(inputs=input_text, outputs=output_layer)

    model = create_model()
    # 优化器也要在策略范围内创建
    optimizer = optimization.create_optimizer(
        init_lr=INIT_LR,
        num_train_steps=num_train_steps,
        num_warmup_steps=num_warmup_steps,
        optimizer_type="adamw"
    )
    model.compile(optimizer=optimizer, loss=loss_function, metrics=[metrics])

3. 修正model.fit()的参数

数据集已经完成批处理，不需要在fit中重复设置batch_size，否则会导致TPU处理冲突：

history = model.fit(train_dataset, epochs=EPOCHS, callbacks=[early_stop, cp_callback])

4. 检查TPU初始化代码

确认TPU初始化步骤完整正确：

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

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