#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2021 The TensorFlow Authors.
End of explanation
"""


import tensorflow as tf
import tensorflow.compat.v1 as tf1

features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [[0.8], [0.9], [1.]]

"""
Explanation: Migrate from Estimator to Keras APIs
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/guide/migrate/migrating_estimator">
    <img src="https://www.tensorflow.org/images/tf_logo_32px.png" />
    View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/migrating_estimator.ipynb">
    <img src="https://www.tensorflow.org/images/colab_logo_32px.png" />
    Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/migrating_estimator.ipynb">
    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
    View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/migrating_estimator.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>

This guide demonstrates how to migrate from TensorFlow 1's tf.estimator.Estimator APIs to TensorFlow 2's tf.keras APIs. First, you will set up and run a basic model for training and evaluation with tf.estimator.Estimator. Then, you will perform the equivalent steps in TensorFlow 2 with the tf.keras APIs. You will also learn how to customize the training step by subclassing tf.keras.Model and using tf.GradientTape.

In TensorFlow 1, the high-level tf.estimator.Estimator APIs let you train and evaluate a model, as well as perform inference and save your model (for serving).
In TensorFlow 2, use the Keras APIs to perform the aforementioned tasks, such as model building, gradient application, training, evaluation, and prediction.

(For migrating model/checkpoint saving workflows to TensorFlow 2, check out the SavedModel and Checkpoint migration guides.)
Setup
Start with imports and a simple dataset:
End of explanation
"""


def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

def _eval_input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

def _model_fn(features, labels, mode):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

"""
Explanation: TensorFlow 1: Train and evaluate with tf.estimator.Estimator
This example shows how to perform training and evaluation with tf.estimator.Estimator in TensorFlow 1.
Start by defining a few functions: an input function for the training data, an evaluation input function for the evaluation data, and a model function that tells the Estimator how the training op is defined with the features and labels:
End of explanation
"""


estimator = tf1.estimator.Estimator(model_fn=_model_fn)
estimator.train(_input_fn)

"""
Explanation: Instantiate your Estimator, and train the model:
End of explanation
"""


estimator.evaluate(_eval_input_fn)

"""
Explanation: Evaluate the program with the evaluation set:
End of explanation
"""


dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss="mse")

"""
Explanation: TensorFlow 2: Train and evaluate with the built-in Keras methods
This example demonstrates how to perform training and evaluation with Keras Model.fit and Model.evaluate in TensorFlow 2. (You can learn more in the Training and evaluation with the built-in methods guide.)

Start by preparing the dataset pipeline with the tf.data.Dataset APIs.
Define a simple Keras Sequential model with one linear (tf.keras.layers.Dense) layer.
Instantiate an Adagrad optimizer (tf.keras.optimizers.Adagrad).
Configure the model for training by passing the optimizer variable and the mean-squared error ("mse") loss to Model.compile.
End of explanation
"""


model.fit(dataset)

"""
Explanation: With that, you are ready to train the model by calling Model.fit:
End of explanation
"""


model.evaluate(eval_dataset, return_dict=True)

"""
Explanation: Finally, evaluate the model with Model.evaluate:
End of explanation
"""


class CustomModel(tf.keras.Sequential):
  """A custom sequential model that overrides `Model.train_step`."""

  def train_step(self, data):
    batch_data, labels = data

    with tf.GradientTape() as tape:
      predictions = self(batch_data, training=True)
      # Compute the loss value (the loss function is configured
      # in `Model.compile`).
      loss = self.compiled_loss(labels, predictions)

    # Compute the gradients of the parameters with respect to the loss.
    gradients = tape.gradient(loss, self.trainable_variables)
    # Perform gradient descent by updating the weights/parameters.
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
    # Update the metrics (includes the metric that tracks the loss).
    self.compiled_metrics.update_state(labels, predictions)
    # Return a dict mapping metric names to the current values.
    return {m.name: m.result() for m in self.metrics}

"""
Explanation: TensorFlow 2: Train and evaluate with a custom training step and built-in Keras methods
In TensorFlow 2, you can also write your own custom training step function with tf.GradientTape to perform forward and backward passes, while still taking advantage of the built-in training support, such as tf.keras.callbacks.Callback and tf.distribute.Strategy. (Learn more in Customizing what happens in Model.fit and Writing custom training loops from scratch.)
In this example, start by creating a custom tf.keras.Model by subclassing tf.keras.Sequential that overrides Model.train_step. (Learn more about subclassing tf.keras.Model). Inside that class, define a custom train_step function that for each batch of data performs a forward pass and backward pass during one training step.
End of explanation
"""


dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

model = CustomModel([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss="mse")

"""
Explanation: Next, as before:
- Prepare the dataset pipeline with tf.data.Dataset.
- Define a simple model with one tf.keras.layers.Dense layer.
- Instantiate Adagrad (tf.keras.optimizers.Adagrad)
- Configure the model for training with Model.compile, while using mean-squared error ("mse") as the loss function.
End of explanation
"""


model.fit(dataset)

"""
Explanation: Call Model.fit to train the model:
End of explanation
"""


model.evaluate(eval_dataset, return_dict=True)

"""
Explanation: And, finally, evaluate the program with Model.evaluate:
End of explanation
"""


