2

I am new to TFX and I would like to know if it's possible to create a TFX pipeline that can train Keras models on TPUs nodes using TPUStrategy from a TPUClusterResolver. Looking at TFX documentation, it is not clear to me if such is possible.

Any feedback about this would be much appreciated! Thank you!

Note: Im using TF v2

TheDude
  • 51
  • 1
  • 4

1 Answers1

0

Yes, TFX Pipeline supports TPU out-of-the-box for TPU-enabled Estimators, the same model_config.model can be used to configure an estimator regardless of whether it is trained on CPU or TPU.

To train in TPU, you should follow the following actions:

  1. Configure a tpu_footprint in your deployment file and the desired trainer stanza of your service file.
  2. Include an empty TPUSettings message in your custom_config file.
  3. And as an optional action, configure a TPUEstimator-specific args in model_config.tpu_estimator.
Eduardo Ortiz
  • 715
  • 3
  • 14