使用 TPU 训练 Keras 模型时出现 InternalError

Question

我正在尝试使用此 link.

在 Tensorflow Hub 的 Google Colab 上微调 BERT 模型

但是，我运行进入如下错误：

InternalError: RET_CHECK failure (third_party/tensorflow/core/tpu/graph_rewrite/distributed_tpu_rewrite_pass.cc:2047) arg_shape.handle_type != DT_INVALID  input edge: [id=2693 model_preprocessing_67660:0 -> cluster_train_function:628]

当我运行我的 model.fit(...) 函数时。

此错误仅在我尝试使用 TPU 时出现（运行在 CPU 上很好，但训练时间很长）。

这是我设置 TPU 和模型的代码：

TPU 设置：

import os
os.environ["TFHUB_MODEL_LOAD_FORMAT"]="UNCOMPRESSED"

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

模型设置：

def build_classifier_model():
  text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')
  preprocessing_layer = hub.KerasLayer('https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3', name='preprocessing')
  encoder_inputs = preprocessing_layer(text_input)
  encoder = hub.KerasLayer('https://tfhub.dev/google/experts/bert/wiki_books/sst2/2', trainable=True, name='BERT_encoder')
  outputs = encoder(encoder_inputs)
  net = outputs['pooled_output']
  net = tf.keras.layers.Dropout(0.1)(net)
  net = tf.keras.layers.Dense(1, activation=None, name='classifier')(net)
  return tf.keras.Model(text_input, net)

模型训练

with strategy.scope():

  bert_model = build_classifier_model()
  loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
  metrics = tf.metrics.BinaryAccuracy()
  epochs = 1
  steps_per_epoch = 1280000
  num_train_steps = steps_per_epoch * epochs
  num_warmup_steps = int(0.1*num_train_steps)

  init_lr = 3e-5
  optimizer = optimization.create_optimizer(init_lr=init_lr,
                                          num_train_steps=num_train_steps,
                                          num_warmup_steps=num_warmup_steps,
                                          optimizer_type='adamw')
  bert_model.compile(optimizer=optimizer,
                         loss=loss,
                         metrics=metrics)
  print(f'Training model')
  history = bert_model.fit(x=X_train, y=y_train,
                               validation_data=(X_val, y_val),
                               epochs=epochs)

请注意，X_train 是一个类型为 str 且形状为 (1280000,) 的 numpy 数组，而 y_train 是一个形状为 (1280000, 1)[=21 的 numpy 数组=]

Answer 1

因为我不完全知道您对代码进行了哪些更改...我不了解您的数据集。但是我可以看到您正在尝试用一个时期训练整个数据集并直接传递每个时期的步骤。我建议这样写

set some batch_size 2^n power (for example 16 or 32 or etc) if you don't want to batch the dataset just set batch_size to 1

batch_size = 16
steps_per_epoch = training_data_size // batch_size

代码的问题很可能是训练数据集的大小。我认为您手动传递训练数据集的值是错误的。

如果您从 tfds 加载数据集，请使用（如 link 所示）：

train_dataset, train_data_size = load_dataset_from_tfds(
  in_memory_ds, tfds_info, train_split, batch_size, bert_preprocess_model)

如果您使用的是自定义数据集，请将已清理数据集的大小放入变量中，然后使用该变量来使用训练数据的大小。尽量避免在代码中手动赋值。

使用 TPU 训练 Keras 模型时出现 InternalError

InternalError when using TPU for training Keras model

python

tensorflow

machine-learning

bert-language-model

tpu