<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Training Keras models with TensorFlow Cloud
Author
Step1:  API overview
Step2:  Let's save the TensorBoard logs and model checkpoints generated during training
in our cloud storage bucket.
Step3:  Here, we will load our data from Keras directly. In general, it's best practice
to store your dataset in your Cloud Storage bucket, however TensorFlow Cloud can
also accomodate datasets stored locally. That's covered in the Multi-file section
of this guide.
Step4:  The TensorFlow Cloud API provides the
remote() function to determine whether code is being executed locally or on
the cloud. This allows for the separate designation of fit() parameters for
local and remote execution, and provides means for easy debugging without overloading
your local machine.
Step5:  Let's save the model in GCS after the training is complete.
Step6:  We can also use this storage bucket for Docker image building, instead of your local
Docker instance. For this, just add your bucket to the docker_image_bucket_name parameter.
Step7:  After training the model, we can load the saved model and view our TensorBoard logs
to monitor performance.
<ASSISTANT_TASK:>
Python Code:
!pip install -q tensorflow_cloud
import tensorflow as tf
import tensorflow_cloud as tfc
from tensorflow import keras
from tensorflow.keras import layers
Explanation: Training Keras models with TensorFlow Cloud
Author: Jonah Kohn<br>
Date created: 2020/08/11<br>
Last modified: 2020/08/11<br>
Description: In-depth usage guide for TensorFlow Cloud.
Introduction
TensorFlow Cloud is a Python package that
provides APIs for a seamless transition from local debugging to distributed training
in Google Cloud. It simplifies the process of training TensorFlow models on the
cloud into a single, simple function call, requiring minimal setup and no changes
to your model. TensorFlow Cloud handles cloud-specific tasks such as creating VM
instances and distribution strategies for your models automatically. This guide
will demonstrate how to interface with Google Cloud through TensorFlow Cloud,
and the wide range of functionality provided within TensorFlow Cloud. We'll start
with the simplest use-case.
Setup
We'll get started by installing TensorFlow Cloud, and importing the packages we
will need in this guide.
End of explanation
def create_model():
    model = keras.Sequential(
        [
            keras.Input(shape=(28, 28)),
            layers.experimental.preprocessing.Rescaling(1.0 / 255),
            layers.Reshape(target_shape=(28, 28, 1)),
            layers.Conv2D(32, 3, activation="relu"),
            layers.MaxPooling2D(2),
            layers.Conv2D(32, 3, activation="relu"),
            layers.MaxPooling2D(2),
            layers.Conv2D(32, 3, activation="relu"),
            layers.Flatten(),
            layers.Dense(128, activation="relu"),
            layers.Dense(10),
        ]
    )
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=keras.metrics.SparseCategoricalAccuracy(),
    )
    return model
Explanation: API overview: a first end-to-end example
Let's begin with a Keras model training script, such as the following CNN:
```python
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
model = keras.Sequential(
    [
        keras.Input(shape=(28, 28)),
        # Use a Rescaling layer to make sure input values are in the [0, 1] range.
        layers.experimental.preprocessing.Rescaling(1.0 / 255),
        # The original images have shape (28, 28), so we reshape them to (28, 28, 1)
        layers.Reshape(target_shape=(28, 28, 1)),
        # Follow-up with a classic small convnet
        layers.Conv2D(32, 3, activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, activation="relu"),
        layers.Flatten(),
        layers.Dense(128, activation="relu"),
        layers.Dense(10),
    ]
)
model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=keras.metrics.SparseCategoricalAccuracy(),
)
model.fit(x_train, y_train, epochs=20, batch_size=128, validation_split=0.1)
```
To train this model on Google Cloud we just need to add a call to run() at
the beginning of the script, before the imports:
python
tfc.run()
You don’t need to worry about cloud-specific tasks such as creating VM instances
and distribution strategies when using TensorFlow Cloud.
The API includes intelligent defaults for all the parameters -- everything is
configurable, but many models can rely on these defaults.
Upon calling run(), TensorFlow Cloud will:
Make your Python script or notebook distribution-ready.
Convert it into a Docker image with required dependencies.
Run the training job on a GCP GPU-powered VM.
Stream relevant logs and job information.
The default VM configuration is 1 chief and 0 workers with 8 CPU cores and
1 Tesla T4 GPU.
Google Cloud configuration
In order to facilitate the proper pathways for Cloud training, you will need to
do some first-time setup. If you're a new Google Cloud user, there are a few
preliminary steps you will need to take:
Create a GCP Project;
Enable AI Platform Services;
Create a Service Account;
Download an authorization key;
Create a Cloud Storage bucket.
Detailed first-time setup instructions can be found in the
TensorFlow Cloud README,
and an additional setup example is shown on the
TensorFlow Blog.
Common workflows and Cloud storage
In most cases, you'll want to retrieve your model after training on Google Cloud.
For this, it's crucial to redirect saving and loading to Cloud Storage while
training remotely. We can direct TensorFlow Cloud to our Cloud Storage bucket for
a variety of tasks. The storage bucket can be used to save and load large training
datasets, store callback logs or model weights, and save trained model files.
To begin, let's configure fit() to save the model to a Cloud Storage, and set
up TensorBoard monitoring to track training progress.
End of explanation
import datetime
import os
# Note: Please change the gcp_bucket to your bucket name.
gcp_bucket = "keras-examples"
checkpoint_path = os.path.join("gs://", gcp_bucket, "mnist_example", "save_at_{epoch}")
tensorboard_path = os.path.join(  # Timestamp included to enable timeseries graphs
    "gs://", gcp_bucket, "logs", datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
)
callbacks = [
    # TensorBoard will store logs for each epoch and graph performance for us.
    keras.callbacks.TensorBoard(log_dir=tensorboard_path, histogram_freq=1),
    # ModelCheckpoint will save models after each epoch for retrieval later.
    keras.callbacks.ModelCheckpoint(checkpoint_path),
    # EarlyStopping will terminate training when val_loss ceases to improve.
    keras.callbacks.EarlyStopping(monitor="val_loss", patience=3),
]
model = create_model()
Explanation: Let's save the TensorBoard logs and model checkpoints generated during training
in our cloud storage bucket.
End of explanation
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
Explanation: Here, we will load our data from Keras directly. In general, it's best practice
to store your dataset in your Cloud Storage bucket, however TensorFlow Cloud can
also accomodate datasets stored locally. That's covered in the Multi-file section
of this guide.
End of explanation
if tfc.remote():
    epochs = 100
    callbacks = callbacks
    batch_size = 128
else:
    epochs = 5
    batch_size = 64
    callbacks = None
model.fit(x_train, y_train, epochs=epochs, callbacks=callbacks, batch_size=batch_size)
Explanation: The TensorFlow Cloud API provides the
remote() function to determine whether code is being executed locally or on
the cloud. This allows for the separate designation of fit() parameters for
local and remote execution, and provides means for easy debugging without overloading
your local machine.
End of explanation
save_path = os.path.join("gs://", gcp_bucket, "mnist_example")
if tfc.remote():
    model.save(save_path)
Explanation: Let's save the model in GCS after the training is complete.
End of explanation
# docs_infra: no_execute
tfc.run(docker_image_bucket_name=gcp_bucket)
Explanation: We can also use this storage bucket for Docker image building, instead of your local
Docker instance. For this, just add your bucket to the docker_image_bucket_name parameter.
End of explanation
# docs_infra: no_execute
model = keras.models.load_model(save_path)
!#docs_infra: no_execute
!tensorboard dev upload --logdir "gs://keras-examples-jonah/logs/fit" --name "Guide MNIST"
Explanation: After training the model, we can load the saved model and view our TensorBoard logs
to monitor performance.
End of explanation
<END_TASK>