python - 使用 TPU pod 训练 keras 模型？

Question

我想知道是否有人有在 TPU 吊舱上使用 keras 模型的示例？

我有一个模型创建方法，它返回一个在 TPU 策略范围内编译的 keras 模型，正如许多将 TPU 与 keras 一起使用的示例所推荐的那样。这适用于 v3-8，但在尝试使用更多内核（特别是 v3-32）时会出错：

with strategy.scope():
    keras_model = create_model()
    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
    keras_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])

运行 model.fit 时，它失败并出现以下错误：

Failed copying input tensor from /job:worker/replica:0/task:0/device:CPU:0 to /job:worker/replica:0/task:1/device:CPU:0 in order to run DatasetFromGraph: FetchOutputs node : not found [Op:DatasetFromGraph]

模型输入采用 numpy 数组的形式。可能需要 tensorflow.data.Dataset 吗？

标签： pythontensorflowkerastpu