<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:  使用分布策略保存和加载模型
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https
Step2:  使用 tf.distribute.Strategy 准备数据和模型：
Step3:  训练模型：
Step4:  保存和加载模型
现在，您已经有一个简单的模型可供使用，让我们了解一下如何保存/加载 API。有两组可用的 API：
高级 Keras model.save 和 tf.keras.models.load_model
低级 tf.saved_model.save 和 tf.saved_model.load
Keras API
以下为使用 Keras API 保存和加载模型的示例：
Step5:  恢复无 tf.distribute.Strategy 的模型：
Step6:  恢复模型后，您可以继续在它上面进行训练，甚至无需再次调用 compile()，因为在保存之前已经对其进行了编译。模型以 TensorFlow 的标准 SavedModel proto 格式保存。有关更多信息，请参阅 saved_model 格式指南。
现在，加载模型并使用 tf.distribute.Strategy 进行训练：
Step7:  如您所见， tf.distribute.Strategy 可以按预期进行加载。此处使用的策略不必与保存前所用策略相同。 
tf.saved_model API
现在，让我们看一下较低级别的 API。保存模型与 Keras API 类似：
Step8:  可以使用 tf.saved_model.load() 进行加载。但是，由于该 API 级别较低（因此用例范围更广泛），所以不会返回 Keras 模型。相反，它返回一个对象，其中包含可用于进行推断的函数。例如：
Step9:  加载的对象可能包含多个函数，每个函数与一个键关联。"serving_default" 是使用已保存的 Keras 模型的推断函数的默认键。要使用此函数进行推断，请运行以下代码：
Step10:  您还可以采用分布式方式加载和进行推断：
Step11:  调用已恢复的函数只是基于已保存模型的前向传递（预测）。如果您想继续训练加载的函数，或者将加载的函数嵌入到更大的模型中，应如何操作？ 通常的做法是将此加载对象包装到 Keras 层以实现此目的。幸运的是，TF Hub 为此提供了 hub.KerasLayer，如下所示：
Step12:  如您所见，hub.KerasLayer 可将从 tf.saved_model.load() 加载回的结果封装到可用于构建其他模型的 Keras 层。这对于迁移学习非常实用。 
我应使用哪种 API？
对于保存，如果您使用的是 Keras 模型，那么始终建议使用 Keras 的 model.save() API。如果您保存的不是 Keras 模型，那么您只能选择使用较低级的 API。
对于加载，使用哪种 API 取决于您要从加载的 API 中获得什么。如果您无法或不想获取 Keras 模型，请使用 tf.saved_model.load()。否则，请使用 tf.keras.models.load_model()。请注意，只有保存 Keras 模型后，才能恢复 Keras 模型。
可以混合使用 API。您可以使用 model.save 保存 Keras 模型，并使用低级 API tf.saved_model.load 加载非 Keras 模型。
Step13:  从本地设备保存/加载
要在远程运行（例如使用 Cloud TPU）的情况下从本地 I/O 设备保存和加载，则必须使用选项 experimental_io_device 将 I/O 设备设置为本地主机。
Step14:  警告
有一种特殊情况，您的 Keras 模型没有明确定义的输入。例如，可以创建没有任何输入形状的序贯模型 (Sequential([Dense(3), ...])。子类化模型在初始化后也没有明确定义的输入。在这种情况下，在保存和加载时都应坚持使用较低级别的 API，否则会出现错误。
要检查您的模型是否具有明确定义的输入，只需检查 model.inputs 是否为 None。如果非 None，则一切正常。在 .fit、.evaluate、.predict 中使用模型，或调用模型 (model(inputs)) 时，输入形状将自动定义。
以下为示例：
<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
import tensorflow_datasets as tfds
import tensorflow as tf
Explanation: 使用分布策略保存和加载模型
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://tensorflow.google.cn/tutorials/distribute/save_and_load"><img src="https://tensorflow.google.cn/images/tf_logo_32px.png">在 TensorFlow.org 上查看</a></td>
  <td><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/distribute/save_and_load.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/save_and_load.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/save_and_load.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png">下载笔记本</a></td>
</table>
概述
在训练期间一般需要保存和加载模型。有两组用于保存和加载 Keras 模型的 API：高级 API 和低级 API。本教程演示了在使用 tf.distribute.Strategy 时如何使用 SavedModel API。要了解 SavedModel 和序列化的相关概况，请参阅保存的模型指南和 Keras 模型序列化指南。让我们从一个简单的示例开始： 
导入依赖项：
End of explanation
mirrored_strategy = tf.distribute.MirroredStrategy()
def get_data():
  datasets, ds_info = tfds.load(name='mnist', with_info=True, as_supervised=True)
  mnist_train, mnist_test = datasets['train'], datasets['test']
  BUFFER_SIZE = 10000
  BATCH_SIZE_PER_REPLICA = 64
  BATCH_SIZE = BATCH_SIZE_PER_REPLICA * mirrored_strategy.num_replicas_in_sync
  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255
    return image, label
  train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
  eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)
  return train_dataset, eval_dataset
def get_model():
  with mirrored_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=[tf.metrics.SparseCategoricalAccuracy()])
    return model
Explanation: 使用 tf.distribute.Strategy 准备数据和模型：
End of explanation
model = get_model()
train_dataset, eval_dataset = get_data()
model.fit(train_dataset, epochs=2)
Explanation: 训练模型：
End of explanation
keras_model_path = "/tmp/keras_save"
model.save(keras_model_path)
Explanation: 保存和加载模型
现在，您已经有一个简单的模型可供使用，让我们了解一下如何保存/加载 API。有两组可用的 API：
高级 Keras model.save 和 tf.keras.models.load_model
低级 tf.saved_model.save 和 tf.saved_model.load
Keras API
以下为使用 Keras API 保存和加载模型的示例：
End of explanation
restored_keras_model = tf.keras.models.load_model(keras_model_path)
restored_keras_model.fit(train_dataset, epochs=2)
Explanation: 恢复无 tf.distribute.Strategy 的模型：
End of explanation
another_strategy = tf.distribute.OneDeviceStrategy("/cpu:0")
with another_strategy.scope():
  restored_keras_model_ds = tf.keras.models.load_model(keras_model_path)
  restored_keras_model_ds.fit(train_dataset, epochs=2)
Explanation: 恢复模型后，您可以继续在它上面进行训练，甚至无需再次调用 compile()，因为在保存之前已经对其进行了编译。模型以 TensorFlow 的标准 SavedModel proto 格式保存。有关更多信息，请参阅 saved_model 格式指南。
现在，加载模型并使用 tf.distribute.Strategy 进行训练：
End of explanation
model = get_model()  # get a fresh model
saved_model_path = "/tmp/tf_save"
tf.saved_model.save(model, saved_model_path)
Explanation: 如您所见， tf.distribute.Strategy 可以按预期进行加载。此处使用的策略不必与保存前所用策略相同。 
tf.saved_model API
现在，让我们看一下较低级别的 API。保存模型与 Keras API 类似：
End of explanation
DEFAULT_FUNCTION_KEY = "serving_default"
loaded = tf.saved_model.load(saved_model_path)
inference_func = loaded.signatures[DEFAULT_FUNCTION_KEY]
Explanation: 可以使用 tf.saved_model.load() 进行加载。但是，由于该 API 级别较低（因此用例范围更广泛），所以不会返回 Keras 模型。相反，它返回一个对象，其中包含可用于进行推断的函数。例如：
End of explanation
predict_dataset = eval_dataset.map(lambda image, label: image)
for batch in predict_dataset.take(1):
  print(inference_func(batch))
Explanation: 加载的对象可能包含多个函数，每个函数与一个键关联。"serving_default" 是使用已保存的 Keras 模型的推断函数的默认键。要使用此函数进行推断，请运行以下代码：
End of explanation
another_strategy = tf.distribute.MirroredStrategy()
with another_strategy.scope():
  loaded = tf.saved_model.load(saved_model_path)
  inference_func = loaded.signatures[DEFAULT_FUNCTION_KEY]
  dist_predict_dataset = another_strategy.experimental_distribute_dataset(
      predict_dataset)
  # Calling the function in a distributed manner
  for batch in dist_predict_dataset:
    another_strategy.run(inference_func,args=(batch,))
Explanation: 您还可以采用分布式方式加载和进行推断：
End of explanation
import tensorflow_hub as hub
def build_model(loaded):
  x = tf.keras.layers.Input(shape=(28, 28, 1), name='input_x')
  # Wrap what's loaded to a KerasLayer
  keras_layer = hub.KerasLayer(loaded, trainable=True)(x)
  model = tf.keras.Model(x, keras_layer)
  return model
another_strategy = tf.distribute.MirroredStrategy()
with another_strategy.scope():
  loaded = tf.saved_model.load(saved_model_path)
  model = build_model(loaded)
  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                optimizer=tf.keras.optimizers.Adam(),
                metrics=[tf.metrics.SparseCategoricalAccuracy()])
  model.fit(train_dataset, epochs=2)
Explanation: 调用已恢复的函数只是基于已保存模型的前向传递（预测）。如果您想继续训练加载的函数，或者将加载的函数嵌入到更大的模型中，应如何操作？ 通常的做法是将此加载对象包装到 Keras 层以实现此目的。幸运的是，TF Hub 为此提供了 hub.KerasLayer，如下所示：
End of explanation
model = get_model()
# Saving the model using Keras's save() API
model.save(keras_model_path) 
another_strategy = tf.distribute.MirroredStrategy()
# Loading the model using lower level API
with another_strategy.scope():
  loaded = tf.saved_model.load(keras_model_path)
Explanation: 如您所见，hub.KerasLayer 可将从 tf.saved_model.load() 加载回的结果封装到可用于构建其他模型的 Keras 层。这对于迁移学习非常实用。 
我应使用哪种 API？
对于保存，如果您使用的是 Keras 模型，那么始终建议使用 Keras 的 model.save() API。如果您保存的不是 Keras 模型，那么您只能选择使用较低级的 API。
对于加载，使用哪种 API 取决于您要从加载的 API 中获得什么。如果您无法或不想获取 Keras 模型，请使用 tf.saved_model.load()。否则，请使用 tf.keras.models.load_model()。请注意，只有保存 Keras 模型后，才能恢复 Keras 模型。
可以混合使用 API。您可以使用 model.save 保存 Keras 模型，并使用低级 API tf.saved_model.load 加载非 Keras 模型。
End of explanation
model = get_model()
# Saving the model to a path on localhost.
saved_model_path = "/tmp/tf_save"
save_options = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
model.save(saved_model_path, options=save_options)
# Loading the model from a path on localhost.
another_strategy = tf.distribute.MirroredStrategy()
with another_strategy.scope():
  load_options = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')
  loaded = tf.keras.models.load_model(saved_model_path, options=load_options)
Explanation: 从本地设备保存/加载
要在远程运行（例如使用 Cloud TPU）的情况下从本地 I/O 设备保存和加载，则必须使用选项 experimental_io_device 将 I/O 设备设置为本地主机。
End of explanation
class SubclassedModel(tf.keras.Model):
  output_name = 'output_layer'
  def __init__(self):
    super(SubclassedModel, self).__init__()
    self._dense_layer = tf.keras.layers.Dense(
        5, dtype=tf.dtypes.float32, name=self.output_name)
  def call(self, inputs):
    return self._dense_layer(inputs)
my_model = SubclassedModel()
# my_model.save(keras_model_path)  # ERROR! 
tf.saved_model.save(my_model, saved_model_path)
Explanation: 警告
有一种特殊情况，您的 Keras 模型没有明确定义的输入。例如，可以创建没有任何输入形状的序贯模型 (Sequential([Dense(3), ...])。子类化模型在初始化后也没有明确定义的输入。在这种情况下，在保存和加载时都应坚持使用较低级别的 API，否则会出现错误。
要检查您的模型是否具有明确定义的输入，只需检查 model.inputs 是否为 None。如果非 None，则一切正常。在 .fit、.evaluate、.predict 中使用模型，或调用模型 (model(inputs)) 时，输入形状将自动定义。
以下为示例：
End of explanation
<END_TASK>