在Colab与Kaggle中训练TensorFlow BERT模型时Notebook卡顿崩溃

代码片段 #

以下是Kaggle Notebook中的相关代码：

数据划分 #

from sklearn.model_selection import train_test_split

X = df[["prompt_text", "text"]].to_numpy()
Y = df[["content", "wording"]].to_numpy()
assert X.shape==Y.shape, "X and Y are not of same shape"

X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=3454)

print("Training data shape:", X_train.shape)
print("Testing data shape:", X_test.shape)

参数设置 #

# Settings
SEQ_LEN = 512
BATCH_SIZE = 32

BERT预处理 #

bert_preprocessor = hub.load("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")
def preprocess(X):
    if (X.ndim==1): seqs = [bert_preprocessor.tokenize(X)]
    else: seqs = [bert_preprocessor.tokenize(X[:,i]) for i in range(X.shape[-1])]    
    return bert_preprocessor.bert_pack_inputs(seqs, seq_length=SEQ_LEN)
bert_train = preprocess(X_train)
bert_test = preprocess(X_test)

构建数据集 #

ds_train = tf.data.Dataset.from_tensor_slices((bert_train, Y_train)) \
            .batch(BATCH_SIZE)
#             .batch(BATCH_SIZE, num_parallel_calls=AUTOTUNE) \
#             .shuffle(ds_train.cardinality(), reshuffle_each_iteration=True).cache()

ds_test = tf.data.Dataset.from_tensor_slices((bert_test, Y_test)) \
            .batch(BATCH_SIZE)
#             .batch(BATCH_SIZE, num_parallel_calls=AUTOTUNE) \
#             .shuffle(ds_test.cardinality(), reshuffle_each_iteration=True).cache()

ds_train, ds_test

模型定义 #

def create_model():
    bert_input = {x: keras.Input(shape=(SEQ_LEN), dtype=tf.int32, name=x)
                  for x in ("input_type_ids", "input_word_ids", "input_mask")}
    bert_outputs = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/4", 
                                  trainable=True, name="BERT")(bert_input)
    drop = keras.layers.Dropout(0.1)(bert_outputs["pooled_output"])
    output = keras.layers.Dense(2)(drop)
    return keras.Model(bert_input, output, name="Summary_Examiner")

模型编译与训练 #

model = create_model()
model.compile(optimizer=keras.optimizers.Adam(),
             loss="mean_squared_error",
             metrics=[keras.metrics.RootMeanSquaredError()])
model.summary()

model.fit(ds_train, epochs=10, validation_data=ds_test, verbose=2)

问题现象 #

在Colab和Kaggle环境中出现两种异常：

• CPU运行时：最后一个训练单元格停留在Epoch 1/10，RAM持续占用直至环境崩溃（Kaggle环境占用超30GB）
• GPU运行时：代码在bert_train = preprocess(X_train)步骤无限卡住，无任何输出

注：数据集总大小仅30MB（训练集+测试集）

已尝试方案 #

• 移除数据集的并行处理与缓存逻辑（代码中已注释）
• 尝试直接使用原始数据训练，不包装为Dataset对象
• 将BATCH_SIZE降至1
• 从fit方法中移除验证数据
• 本地Windows机器无法运行（TensorFlow-text不支持该环境）

恳请各位提供解决思路，已尝试所有能想到的方法但问题仍未解决。

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