<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2019 The TensorFlow Authors.
Step1:  Keras 的分布式训练
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  下载数据集
下载 MNIST 数据集并从 TensorFlow Datasets 加载。 这会返回 tf.data 格式的数据集。
将 with_info 设置为 True 会包含整个数据集的元数据,其中这些数据集将保存在 info 中。 除此之外，该元数据对象包括训练和测试示例的数量。
Step3:  定义分配策略
创建一个 MirroredStrategy 对象。这将处理分配策略，并提供一个上下文管理器（tf.distribute.MirroredStrategy.scope）来构建你的模型。
Step4:  设置输入管道（pipeline）
在训练具有多个 GPU 的模型时，您可以通过增加批量大小（batch size）来有效地使用额外的计算能力。通常来说，使用适合 GPU 内存的最大批量大小（batch size），并相应地调整学习速率。
Step5:  0-255 的像素值， 必须标准化到 0-1 范围。在函数中定义标准化。
Step6:  将此功能应用于训练和测试数据，随机打乱训练数据，并批量训练。 请注意，我们还保留了训练数据的内存缓存以提高性能。
Step7:  生成模型
在 strategy.scope 的上下文中创建和编译 Keras 模型。
Step8:  定义回调（callback）
这里使用的回调（callbacks）是：
TensorBoard
Step9:  训练和评估
在该部分，以普通的方式训练模型，在模型上调用 fit 并传入在教程开始时创建的数据集。 无论您是否分布式训练，此步骤都是相同的。
Step10:  如下所示，检查点（checkpoint）将被保存。
Step11:  要查看模型的执行方式，请加载最新的检查点（checkpoint）并在测试数据上调用 evaluate 。
使用适当的数据集调用 evaluate 。
Step12:  要查看输出，您可以在终端下载并查看 TensorBoard 日志。
$ tensorboard --logdir=path/to/log-directory
Step13:  导出到 SavedModel
将图形和变量导出为与平台无关的 SavedModel 格式。 保存模型后，可以在有或没有 scope 的情况下加载模型。
Step14:  在无需 strategy.scope 加载模型。
Step15:  在含 strategy.scope 加载模型。
<ASSISTANT_TASK:>
Python Code:
#@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 2019 The TensorFlow Authors.
End of explanation
# 导入 TensorFlow 和 TensorFlow 数据集
import tensorflow_datasets as tfds
import tensorflow as tf
tfds.disable_progress_bar()
import os
print(tf.__version__)
Explanation: Keras 的分布式训练
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://tensorflow.google.cn/tutorials/distribute/keras"><img src="https://tensorflow.google.cn/images/tf_logo_32px.png" />在 tensorflow.google.cn 上查看</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/distribute/keras.ipynb"><img src="https://tensorflow.google.cn/images/colab_logo_32px.png" />在 Google Colab 运行</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/distribute/keras.ipynb"><img src="https://tensorflow.google.cn/images/GitHub-Mark-32px.png" />在 Github 上查看源代码</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/tutorials/distribute/keras.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png" />下载此 notebook</a>
  </td>
</table>
Note: 我们的 TensorFlow 社区翻译了这些文档。因为社区翻译是尽力而为， 所以无法保证它们是最准确的，并且反映了最新的
官方英文文档。如果您有改进此翻译的建议， 请提交 pull request 到
tensorflow/docs GitHub 仓库。要志愿地撰写或者审核译文，请加入
docs-zh-cn@tensorflow.org Google Group。
概述
tf.distribute.Strategy API 提供了一个抽象的 API ，用于跨多个处理单元（processing units）分布式训练。它的目的是允许用户使用现有模型和训练代码，只需要很少的修改，就可以启用分布式训练。
本教程使用 tf.distribute.MirroredStrategy，这是在一台计算机上的多 GPU（单机多卡）进行同时训练的图形内复制（in-graph replication）。事实上，它会将所有模型的变量复制到每个处理器上，然后，通过使用 all-reduce 去整合所有处理器的梯度（gradients），并将整合的结果应用于所有副本之中。
MirroredStategy 是 tensorflow 中可用的几种分发策略之一。 您可以在 分发策略指南 中阅读更多分发策略。
Keras API
这个例子使用 tf.keras API 去构建和训练模型。 关于自定义训练模型，请参阅 tf.distribute.Strategy with training loops 教程。
导入依赖
End of explanation
datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)
mnist_train, mnist_test = datasets['train'], datasets['test']
Explanation: 下载数据集
下载 MNIST 数据集并从 TensorFlow Datasets 加载。 这会返回 tf.data 格式的数据集。
将 with_info 设置为 True 会包含整个数据集的元数据,其中这些数据集将保存在 info 中。 除此之外，该元数据对象包括训练和测试示例的数量。
End of explanation
strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
Explanation: 定义分配策略
创建一个 MirroredStrategy 对象。这将处理分配策略，并提供一个上下文管理器（tf.distribute.MirroredStrategy.scope）来构建你的模型。
End of explanation
# 您还可以执行 info.splits.total_num_examples 来获取总数
# 数据集中的样例数量。
num_train_examples = info.splits['train'].num_examples
num_test_examples = info.splits['test'].num_examples
BUFFER_SIZE = 10000
BATCH_SIZE_PER_REPLICA = 64
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
Explanation: 设置输入管道（pipeline）
在训练具有多个 GPU 的模型时，您可以通过增加批量大小（batch size）来有效地使用额外的计算能力。通常来说，使用适合 GPU 内存的最大批量大小（batch size），并相应地调整学习速率。
End of explanation
def scale(image, label):
  image = tf.cast(image, tf.float32)
  image /= 255
  return image, label
Explanation: 0-255 的像素值， 必须标准化到 0-1 范围。在函数中定义标准化。
End of explanation
train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)
Explanation: 将此功能应用于训练和测试数据，随机打乱训练数据，并批量训练。 请注意，我们还保留了训练数据的内存缓存以提高性能。
End of explanation
with strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(64, activation='relu'),
      tf.keras.layers.Dense(10)
  ])
  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                optimizer=tf.keras.optimizers.Adam(),
                metrics=['accuracy'])
Explanation: 生成模型
在 strategy.scope 的上下文中创建和编译 Keras 模型。
End of explanation
# 定义检查点（checkpoint）目录以存储检查点（checkpoints）
checkpoint_dir = './training_checkpoints'
# 检查点（checkpoint）文件的名称
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")
# 衰减学习率的函数。
# 您可以定义所需的任何衰减函数。
def decay(epoch):
  if epoch < 3:
    return 1e-3
  elif epoch >= 3 and epoch < 7:
    return 1e-4
  else:
    return 1e-5
# 在每个 epoch 结束时打印LR的回调（callbacks）。
class PrintLR(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    print('\nLearning rate for epoch {} is {}'.format(epoch + 1,
                                                      model.optimizer.lr.numpy()))
callbacks = [
    tf.keras.callbacks.TensorBoard(log_dir='./logs'),
    tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix,
                                       save_weights_only=True),
    tf.keras.callbacks.LearningRateScheduler(decay),
    PrintLR()
]
Explanation: 定义回调（callback）
这里使用的回调（callbacks）是：
TensorBoard: 此回调（callbacks）为 TensorBoard 写入日志，允许您可视化图形。
Model Checkpoint: 此回调（callbacks）在每个 epoch 后保存模型。
Learning Rate Scheduler: 使用此回调（callbacks），您可以安排学习率在每个 epoch/batch 之后更改。
为了便于说明，添加打印回调（callbacks）以在笔记本中显示学习率。
End of explanation
model.fit(train_dataset, epochs=12, callbacks=callbacks)
Explanation: 训练和评估
在该部分，以普通的方式训练模型，在模型上调用 fit 并传入在教程开始时创建的数据集。 无论您是否分布式训练，此步骤都是相同的。
End of explanation
# 检查检查点（checkpoint）目录
!ls {checkpoint_dir}
Explanation: 如下所示，检查点（checkpoint）将被保存。
End of explanation
model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
eval_loss, eval_acc = model.evaluate(eval_dataset)
print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
Explanation: 要查看模型的执行方式，请加载最新的检查点（checkpoint）并在测试数据上调用 evaluate 。
使用适当的数据集调用 evaluate 。
End of explanation
!ls -sh ./logs
Explanation: 要查看输出，您可以在终端下载并查看 TensorBoard 日志。
$ tensorboard --logdir=path/to/log-directory
End of explanation
path = 'saved_model/'
model.save(path, save_format='tf')
Explanation: 导出到 SavedModel
将图形和变量导出为与平台无关的 SavedModel 格式。 保存模型后，可以在有或没有 scope 的情况下加载模型。
End of explanation
unreplicated_model = tf.keras.models.load_model(path)
unreplicated_model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=tf.keras.optimizers.Adam(),
    metrics=['accuracy'])
eval_loss, eval_acc = unreplicated_model.evaluate(eval_dataset)
print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
Explanation: 在无需 strategy.scope 加载模型。
End of explanation
with strategy.scope():
  replicated_model = tf.keras.models.load_model(path)
  replicated_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                           optimizer=tf.keras.optimizers.Adam(),
                           metrics=['accuracy'])
  eval_loss, eval_acc = replicated_model.evaluate(eval_dataset)
  print ('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
Explanation: 在含 strategy.scope 加载模型。
End of explanation
<END_TASK>