Training a keras model using TPU pods?

Question

I was wondering if anyone has an example of using a keras model on a TPU pod?

I have a model creating method which returns a keras model which is compiled within a TPU strategy scope, as recommended by many examples on using TPUs with keras. This works with v3-8 but gives an error when tried with more cores (specifically 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'])

When running model.fit, it fails with the following error:

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]

The model input are in the form of numpy arrays. Is perhaps a tensorflow.data.Dataset required?

Training a keras model using TPU pods?

0 Answers