.. _keras-model-to-cerebras-estimator:

Keras Model to CerebrasEstimator
================================

The ``KerasModelToCerebrasEstimator`` function should be used to convert your Keras model to the ``CerebrasEstimator``. This section describes the ``KerasModelToCerebrasEstimator`` function. For more on ``CerebrasEstimator``, see :ref:`interface-cerebras-estimator`.

.. _use-mixed-precision:

Using the KerasModelToCerebrasEstimator
---------------------------------------
The ``KerasModelToCerebrasEstimator`` is a wrapper that converts the Keras model so the model can be run using the ``CerebrasEstimator``.

.. important::

	Make sure that you use mixed precision by specifying ``mixed_float16`` while using the ``KerasModelToCerebrasEstimator`` wrapper.


Example
~~~~~~~

The following example code shows how you can construct the ``KerasModelToCerebrasEstimator``:

.. code-block:: python

  from cerebras.tf.cs_model_to_estimator import KerasModelToCerebrasEstimator, KerasModelToModelFn

  ...

  dtype = tf.keras.mixed_precision.experimental.Policy(
        'mixed_float16', # Important: This is required.
        loss_scale=None
    )

  tf.keras.mixed_precision.experimental.set_policy(dtype)

  ...

  estimator = KerasModelToCerebrasEstimator(
                 model_fn=model_fn,
                 model_dir=None,
                 compile_dir=None,
                 config: CSRunConfig=None,
                 params=None,
              )

  estimator.compile(
      input_fn=input_fn,
  )


where:

- ``model_fn``: *Function*. Required. The Keras model function.
- ``model_dir``: *String*. Optional. Same as the ``model_dir`` passed to the ``tf.estimator``. The location where your model and all the outputs such as checkpoints, summaries and event files are stored. Default value is ``None``. See also `tf.estimator.Estimator <https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator>`__.
- ``compile_dir``: *String*. Optional. The directory where the compilation results are stored, and from which the compilation outputs are reloaded. Default value is ``model_dir``.
- ``config``: *CSRunConfig*. Optional. A ``CSRunConfig`` object specifying the environment and runtime configuration options. These options are an extension of `TensorFlow RunConfig <https://www.tensorflow.org/api_docs/python/tf/estimator/RunConfig>`__. Default value is ``None``.
- ``params``: *Dictionary*. Optional. A parameters dictionary that contains additional configuration information that will be passed to ``model_fn`` and ``input_fn``. Default value is ``None``.


After converting to the ``CerebrasEstimator``, you can run the compile on the model by calling the ``CerebrasEstimator`` method: ``estimator.compile()``.