save_and_load.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

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

분산 전략을 사용한 모델 저장 및 불러오기

TensorFlow.org에서 보기

Google Colab에서 실행하기

GitHub에서 소스 보기

노트북 다운로드하기

개요

이 튜토리얼에서는 훈련 중 또는 훈련 후에 tf.distribute.Strategy를 사용하여 SavedModel 형식으로 모델을 저장하고 로드하는 방법을 보여줍니다. Keras 모델을 저장하고 로드하는 API에는 상위 수준(tf.keras.Model.save 및 tf.keras.models.load_model) 및 하위 수준(tf.saved_model.save 및 tf.saved_model.load)의 두 가지 종류가 있습니다.

SavedModel 및 직렬화에 대한 일반적인 내용은 저장된 모델 가이드 및 Keras 모델 직렬화 가이드를 참조하세요. 간단한 예부터 시작하겠습니다.

주의: TensorFlow 모델은 코드이며 신뢰할 수 없는 코드에 주의하는 것이 중요합니다. 자세한 내용은 TensorFlow 안전하게 사용하기를 참조하세요.

필요한 패키지 가져오기:

import tensorflow_datasets as tfds

import tensorflow as tf

TensorFlow Datasets 및 tf.data를 사용하여 데이터를 로드 및 준비하고 tf.distribute.MirroredStrategy를 사용하여 모델을 생성합니다.

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을 사용하여 모델을 훈련시킵니다.

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

모델 저장하고 불러오기

이제 작업할 간단한 모델이 생겼으므로 API 저장/로드하기를 살펴보겠습니다. 두 가지 종류의 API를 사용할 수 있습니다.

• 고수준 (Keras): Model.save 및 tf.keras.models.load_model(.keras zip 아카이브 형식)

• 저수준: tf.saved_model.save 및 tf.saved_model.load(TF SavedModel 형식)

Keras API

다음은 Keras API를 사용하여 모델을 저장하고 로드하는 예입니다.

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

tf.distribute.Strategy없이 모델 복원시키기:

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

모델을 복원한 후 Model.compile을 다시 호출할 필요 없이 모델에 대한 훈련을 계속할 수 있습니다. 저장하기 전에 이미 컴파일되었기 때문입니다. 모델은 .keras 확장자로 표시된 Keras zip 아카이브 형식으로 저장됩니다. 자세한 내용은 케라스 저장 가이드를 참조하세요.

이제 모델을 복원하고 tf.distribute.Strategy를 사용하여 훈련시킵니다.

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와 유사합니다.

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(따라서 사용 사례의 범위가 더 넓음)이기 때문에 Keras 모델을 반환하지 않습니다. 대신 추론을 수행하는 데 사용할 수 있는 함수가 포함된 객체를 반환합니다. 예를 들면 다음과 같습니다.

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

로드된 객체에는 각각 키와 연결된 여러 함수가 포함될 수 있습니다. "serving_default" 키는 저장된 Keras 모델이 있는 추론 함수의 기본 키입니다. 이 함수로 추론하려면 다음과 같이 합니다.

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

또한 분산방식으로 불러오고 추론할 수 있습니다:

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가 있습니다.

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)

위의 예에서 Tensorflow Hub의 hub.KerasLayer는 tf.saved_model.load에서 다시 로드된 결과를 다른 모델을 빌드하는 데 사용되는 Keras 레이어로 래핑합니다. 이것은 전이 학습에 매우 유용합니다.

어떤 API를 사용해야 할까요?

저장의 경우, Keras 모델로 작업한다면 하위 수준 API에서 제공하는 추가적 제어가 필요한 경우가 아니면 Model.save API를 사용합니다. 저장하는 대상이 Keras 모델이 아닌 경우 하위 수준 API인 tf.saved_model.save가 유일한 선택입니다.

로드의 경우, API 선택은 모델 로드 API에서 얻고자 하는 결과에 따라 다릅니다. Keras 모델을 가져올 수 없거나 원하지 않으면 tf.saved_model.load를 사용합니다. 그렇지 않으면 tf.keras.models.load_model을 사용합니다. Keras 모델을 저장한 경우에만 Keras 모델을 다시 가져올 수 있습니다.

API를 혼합하여 구성할 수 있습니다. Model.save를 사용하여 Keras 모델을 저장하고 하위 수준 API인 tf.saved_model.load를 사용하여 비 Keras 모델을 로드할 수 있습니다.

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로 설정해야 합니다. 예를 들면 다음과 같습니다.

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 모델을 생성한 다음 훈련 전에 저장할 때입니다. 예를 들면 다음과 같습니다.

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을 찾거나 생성할 수 없기 때문입니다.

주의: 적어도 하나의 ConcreteFunction 없이 모델을 저장하면 안 됩니다. 그렇지 않으면 하위 수준 API가 ConcreteFunction 서명 없이 SavedModel을 생성하기 때문입니다(SavedModel 형식에 대해 자세히 알아보기). 예를 들면 다음과 같습니다.

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

일반적으로 모델의 순방향 전달(call 메서드)은 모델이 처음으로 호출될 때 종종 Model.fit 메서드를 통해 자동으로 추적됩니다. 예를 들어 첫 번째 레이어를 tf.keras.layers.InputLayer 또는 다른 레이어 유형으로 만들고 이를 input_shape 키워드 인수를 전달하여 입력 형상을 설정하면 Keras 순차형 및 함수형 API에서 ConcreteFunction을 생성할 수도 있습니다.

모델에 추적된 ConcreteFunction이 있는지 확인하려면 Model.save_spec이 None인지 확인하세요.

print(my_model.save_spec() is None)

tf.keras.Model.fit을 사용하여 모델을 훈련하고 save_spec이 정의되고 모델 저장이 작동하는지 확인합니다.

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)