Do TFX pipelines support training on TPUs?

TensorFlow General Discussion
, ,
1

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 y

3

Hi @Neural, Yes it is possible to train a model on TPU nodes. Please refer to this document for implementation. In the document in _get_distribution_strategy function replace tf.distribute.MirroredStrategy with tf.distribute.TPUStrategy and pass values to the required arguments by doing so you can train the model on TPU. Thank You.