Distributed training with TensorFlow 2


The managed MLflow integration with Databricks on Google Cloud requires Databricks Runtime for Machine Learning 9.1 LTS or above.

spark-tensorflow-distributor is an open-source native package in TensorFlow that helps users do distributed training with TensorFlow on their Spark clusters. It is built on top of tensorflow.distribute.Strategy, which is one of the major features in TensorFlow 2. For detailed API documentation, see docstrings. For general documentation about distributed TensorFlow, see Distributed training with TensorFlow.

Example notebook

Distributed Training with TensorFlow 2

Open notebook in new tab