GitHub Repository: tensorflow/docs-l10n
Path: blob/master/site/ja/tutorials/distribute/save_and_load.ipynb
²⁵¹¹⁸ views

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

In [ ]:

#@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.

分散ストラテジーを使ってモデルを保存して読み込む

概要

このチュートリアルでは、トレーニング中またはトレーニング後に tf.distribute.Strategy を使用して SavedModel 形式でモデルを保存して読み込む方法を説明します。Keras モデルの保存と読み込みには、高レベル（tf.keras.Model.save と tf.keras.models.load_model）と低レベル（tf.saved_model.save と tf.saved_model.load）の 2 種類の API があります。

SavedModel とシリアル化の全般的な内容については、SavedModel ガイドと Keras モデルのシリアル化ガイドをお読みください。では、単純な例から始めましょう。

注意: TensorFlow モデルはコードであるため、信頼できないコードには注意する必要があります。詳細は、TensorFlow を安全に使用するをご覧ください。

依存関係をインポートします。

In [ ]:

import tensorflow_datasets as tfds

import tensorflow as tf

TensorFlow Datasets と tf.data でデータを読み込んで準備し、tf.distribute.MirroredStrategy を使ってモデルを作成します。

In [ ]:

mirrored_strategy = tf.distribute.MirroredStrategy()

def get_data():
  datasets = tfds.load(name='mnist', 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

tf.keras.Model.fit を使用してモデルをトレーニングします。

In [ ]:

model = get_model()
train_dataset, eval_dataset = get_data()
model.fit(train_dataset, epochs=2)

モデルを保存して読み込む

作業に使用する単純なモデルを準備できたので、保存と読み込みに使用する API を見てみましょう。使用できる API には、以下の 2 種類があります。

高レベル（Keras）: Model.save および tf.keras.models.load_model（.keras zip アーカイブ形式）
低レベル: tf.saved_model.save および tf.saved_model.load（TF SavedModel 形式）

Keras API

Keras API を使用したモデルの保存と読み込みの例を以下に示します。

In [ ]:

keras_model_path = '/tmp/keras_save.keras'
model.save(keras_model_path)

tf.distribute.Strategy を使用せずにモデルを復元します。

In [ ]:

restored_keras_model = tf.keras.models.load_model(keras_model_path)
restored_keras_model.fit(train_dataset, epochs=2)

モデルを復元したら、Model.compile をもう一度呼び出さずにそのままトレーニングを続行できます。これは、保存前にすでにコンパイル済みであるためです。このモデルは、Keras zip アーカイブ形式で保存されており、.keras 拡張子で識別できます。詳細については、Keras の保存に関するガイドをご覧ください。

次に、tf.distribute.Strategy を使用してモデルを復元し、トレーニングします。

In [ ]:

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)

Model.fit 出力からわかるように、tf.distribute.Strategy を使って期待どおり読み込まれました。ここで使用されるストラテジーは、保存前と同じストラテジーである必要はありません。

`tf.saved_model` API

より低レベルの API を使用したモデルの保存方法は、Keras API を使う方法に似ています。

In [ ]:

model = get_model()  # get a fresh model
saved_model_path = '/tmp/tf_save'
tf.saved_model.save(model, saved_model_path)

読み込みは、tf.saved_model.load を使用して行えますが、これは低レベル API（したがって、より幅広いユースケースのある API）であるため、Keras モデルを返しません。代わりに、推論を行うために使用できる関数を含むオブジェクトを返します。以下に例を示します。

In [ ]:

DEFAULT_FUNCTION_KEY = 'serving_default'
loaded = tf.saved_model.load(saved_model_path)
inference_func = loaded.signatures[DEFAULT_FUNCTION_KEY]

読み込まれたオブジェクトには、それぞれにキーが関連付けられた複数の関数が含まれている可能性があります。"serving_default" キーは、保存された Keras モデルを使用した推論関数のデフォルトのキーです。この関数で推論するには、以下のようにします。

In [ ]:

predict_dataset = eval_dataset.map(lambda image, label: image)
for batch in predict_dataset.take(1):
  print(inference_func(batch))

また、分散方法で読み込んで推論を実行することもできます。

In [ ]:

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:
    result = another_strategy.run(inference_func, args=(batch,))
    print(result)
    break

復元された関数の呼び出しは、保存されたモデル（tf.keras.Model.predict）に対するフォワードパスです。読み込まれた関数をトレーニングし続ける場合はどうでしょうか。または読み込まれた関数をより大きなモデルに埋め込むには？一般的には、この読み込まれたオブジェクトを Keras レイヤーにラップして行うことができます。幸いにも、TF Hub には、以下に示すとおり、この目的に使用できる hub.KerasLayer が用意されています。

In [ ]:

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)

上記の例では、hub.KerasLayer は tf.saved_model.load() から読み込まれた結果を、別のモデルの構築に使用できる Keras レイヤーにラップしています。転移学習を行う際に非常に便利な手法です。

どの API を使用すべきですか？

保存においては、Keras モデルを使用している場合は、低レベル API が実現できる追加の制御が必要でない限り、Keras の Model.save API を使用します。保存しているものが Keras モデルでない場合は、低レベル API の tf.saved_model.save しか使用できません。

読み込みにおいては、使用する API はモデルの読み込みから得ようとしているものによって異なります。Keras モデルを使用できない場合（または使用したくない場合）は、tf.saved_model.load を使用し、使用できる場合は tf.keras.models.load_model を使用します。Keras モデルを保存した場合にのみ、Keras モデルを読み込めることに注意してください。

API を混在させることも可能です。model.save で Keras モデルを保存し、低レベルの tf.saved_model.load API を使用して、非 Keras モデルを読み込むことができます。

In [ ]:

model = get_model()

# Saving the model using Keras `Model.save`
model.save(saved_model_path)

another_strategy = tf.distribute.MirroredStrategy()
# Loading the model using the lower-level API
with another_strategy.scope():
  loaded = tf.saved_model.load(saved_model_path)

ローカルデバイスからの読み込みまたは保存

ローカル I/O デバイスから読み込みと保存を行い、リモートデバイスでトレーニングする場合（Cloud TPU を使用する場合など）、tf.saved_model.SaveOptions と tf.saved_model.LoadOptions に experimental_io_device を使用して、I/O デバイスを localhost に設定する必要があります。以下に例を示します。

In [ ]:

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)

警告

Keras モデルを特定の方法で作成してから、トレーニングする前に保存するという、以下のような特別なケースがあります。

In [ ]:

class SubclassedModel(tf.keras.Model):
  """Example model defined by subclassing `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()
try:
  my_model.save(saved_model_path)
except ValueError as e:
  print(f'{type(e).__name__}: ', *e.args)

SavedModel は tf.function をトレースする際に生成される tf.types.experimental.ConcreteFunction オブジェクトを保存します（詳細は、グラフと tf.function の基本ガイドの関数はいつトレースしますか？ をご覧ください）。このような ValueError が発生した場合、Model.save がトレースされた ConcreteFunction を見つけられなかったか作成できなかったことが原因です。

注意: 少なくとも 1 つの ConcreteFunction がない場合にモデルを保存しないことをお勧めします。そうでない場合、低レベル API は、ConcreteFunction シグネチャのない状態で SavedModel を生成してしまうためです（SavedModel 形式については、こちらをご覧ください）。以下に例を示します。

In [ ]:

tf.saved_model.save(my_model, saved_model_path)
x = tf.saved_model.load(saved_model_path)
x.signatures

通常、モデルのフォワードパス（call メソッド）は、モデルが Keras の Model.fit メソッドを通じて初めて呼び出されたときに、自動的にトレースされます。また、最初のレイヤーを tf.keras.layers.InputLayer などにして、input_shape キーワード引数に渡すことで入力形状を設定している場合、Keras の Sequential API と Functional API によって ConcreteFunction が生成されることもあります。

モデルにトレース済みの ConcreteFunction が存在するかを確認するには、Model.save_spec が None になっていることを確認します。

In [ ]:

print(my_model.save_spec() is None)

tf.keras.Model.fit を使ってモデルをトレーニングし、save_spec が定義され、モデルの保存が機能するかを確認しましょう。

In [ ]:

BATCH_SIZE_PER_REPLICA = 4
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * mirrored_strategy.num_replicas_in_sync

dataset_size = 100
dataset = tf.data.Dataset.from_tensors(
    (tf.range(5, dtype=tf.float32), tf.range(5, dtype=tf.float32))
    ).repeat(dataset_size).batch(BATCH_SIZE)

my_model.compile(optimizer='adam', loss='mean_squared_error')
my_model.fit(dataset, epochs=2)

print(my_model.save_spec() is None)
my_model.save(saved_model_path)

Copyright 2019 The TensorFlow Authors.

分散ストラテジーを使ってモデルを保存して読み込む

概要

モデルを保存して読み込む

Keras API

`tf.saved_model` API

どの API を使用すべきですか？

ローカルデバイスからの読み込みまたは保存

警告

Product

Resources

Company

Copyright 2019 The TensorFlow Authors.

分散ストラテジーを使ってモデルを保存して読み込む

概要

モデルを保存して読み込む

Keras API

tf.saved_model API

どの API を使用すべきですか？

ローカルデバイスからの読み込みまたは保存

警告

`tf.saved_model` API