:orphan:

Kubeflow TensorFlow
===================

TensorFlow operator is useful to natively run distributed TensorFlow training jobs on Flyte.
It is a wrapper built around `Kubeflow's TensorFlow operator <https://www.kubeflow.org/docs/components/training/tftraining/>`__.

Installation
------------

To install the Kubeflow TensorFlow plugin, run the following command:

.. code-block:: bash

    pip install flytekitplugins-kftensorflow

To enable the plugin in the backend, follow instructions outlined in the :std:ref:`flyte:deployment-plugin-setup-k8s` guide.

Code
----

We will write an example that does distributed training using the Kubeflow TensorFlow operator.
Before that, let's look at the compute setup and Dockerfile.

GPU to CPU
^^^^^^^^^^

**GPU support has been enabled in the code by default**. If you want to test your code on a CPU, incorporate the following changes:

* Replace ``FROM tensorflow/tensorflow:latest-gpu`` with ``FROM tensorflow/tensorflow:latest`` in the Dockerfile
* Remove the ``gpu`` parameter from the ``Resources`` definition in the example

Dockerfile
^^^^^^^^^^

The example uses TensorFlow-GPU image.

.. literalinclude:: ../../../../../integrations/kubernetes/kftensorflow/Dockerfile
    :emphasize-lines: 1-3
    :language: docker



.. raw:: html

    <div class="sphx-glr-thumbnails">


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="When you need to scale up model training using TensorFlow, you can use :py:class:`~tensorflow:t...">

.. only:: html

  .. image:: /auto/integrations/kubernetes/kftensorflow/images/thumb/sphx_glr_tf_mnist_thumb.png
    :alt: Distributed TensorFlow Training

  :ref:`sphx_glr_auto_integrations_kubernetes_kftensorflow_tf_mnist.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Distributed TensorFlow Training</div>
    </div>


.. raw:: html

    </div>


.. toctree::
   :hidden:

   /auto/integrations/kubernetes/kftensorflow/tf_mnist


.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-gallery

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download all examples in Python source code: kftensorflow_python.zip </auto/integrations/kubernetes/kftensorflow/kftensorflow_python.zip>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download all examples in Jupyter notebooks: kftensorflow_jupyter.zip </auto/integrations/kubernetes/kftensorflow/kftensorflow_jupyter.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_