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.

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

Note: 이 문서는 텐서플로 커뮤니티에서 번역했습니다. 커뮤니티 번역 활동의 특성상 정확한 번역과 최신 내용을 반영하기 위해 노력함에도 불구하고 공식 영문 문서의 내용과 일치하지 않을 수 있습니다. 이 번역에 개선할 부분이 있다면 tensorflow/docs-l10n 깃허브 저장소로 풀 리퀘스트를 보내주시기 바랍니다. 문서 번역이나 리뷰에 참여하려면 docs-ko@tensorflow.org로 메일을 보내주시기 바랍니다.

개요

훈련 도중 모델을 저장하고 불러오는 것은 흔히 일어나는 일입니다. 케라스 모델을 저장하고 불러오기 위한 API에는 high-level API와 low-level API, 두 가지가 있습니다. 이 튜토리얼은 tf.distribute.Strategy를 사용할 때 어떻게 SavedModel APIs를 사용할 수 있는지 보여줍니다. SavedModel과 직렬화에 관한 일반적인 내용을 학습하려면, saved model guideKeras model serialization guide를 읽어보는 것을 권장합니다. 간단한 예로 시작해보겠습니다:

필요한 패키지 가져오기:


In [ ]:
!pip install tf-nightly
import tensorflow_datasets as tfds

import tensorflow as tf
tfds.disable_progress_bar()

tf.distribute.Strategy를 사용하는 모델과 데이터 준비하기:


In [ ]:
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=['accuracy'])
    return model

모델 훈련시키기:


In [ ]:
model = get_model()
train_dataset, eval_dataset = get_data()
model.fit(train_dataset, epochs=2)

모델 저장하고 불러오기

이제 사용할 모델을 가지고 있으므로 API를 이용해 모델을 저장하고 불러오는 방법에 대해 살펴봅시다. 두 가지 API를 사용 할 수 있습니다:

  • 고수준 케라스 model.savetf.keras.models.load_model
  • 저수준 케라스 tf.saved_model.savetf.saved_model.load

케라스 API

케라스 API들을 이용해 모델을 저장하고 불러오는 예를 소개합니다.


In [ ]:
keras_model_path = "/tmp/keras_save"
model.save(keras_model_path)  # save()는 전략 범위를 벗어나 호출되어야 합니다.

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


In [ ]:
restored_keras_model = tf.keras.models.load_model(keras_model_path)
restored_keras_model.fit(train_dataset, epochs=2)

모델을 복원시킨 후에는 compile()이 이미 저장되기 전에 컴파일 되기 때문에, compile()을 다시 호출하지 않고도 모델 훈련을 계속 할 수 있습니다. 그 모델은 텐서플로 표준 SavedModel의 프로토 타입에 저장됩니다. 더 많은 정보를 원한다면, the guide to saved_model format를 참고하세요.

tf.distribute.strategy의 범위를 벗어나서 model.save() 방법을 호출하는 것은 중요합니다. 범위 안에서 호출하는 것은 지원하지 않습니다.

이제 모델을 불러와서 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)

위에서 볼 수 있듯이, 불러오기는 tf.distribute.Strategy에서 예상한대로 작동합니다. 여기서 사용된 전략은 이전에 사용된 전략과 같지 않아도 됩니다.

tf.saved_model 형 API

이제 저수준 API에 대해서 살펴봅시다. 모델을 저장하는 것은 케라스 API와 비슷합니다:


In [ ]:
model = get_model()  # 새 모델 얻기
saved_model_path = "/tmp/tf_save"
tf.saved_model.save(model, saved_model_path)

tf.saved_model.load()로 불러올 수 있습니다. 그러나 저수준 단계의 API이기 때문에 (따라서 더 넓은 사용범위를 갖습니다), 케라스 모델을 반환하지 않습니다. 대신, 추론하기 위해 사용될 수 있는 기능들을 포함한 객체를 반환합니다. 예를 들어:


In [ ]:
DEFAULT_FUNCTION_KEY = "serving_default"
loaded = tf.saved_model.load(saved_model_path)
inference_func = loaded.signatures[DEFAULT_FUNCTION_KEY]

불러와진 객체는 각각 키와 관련된 채, 여러 기능을 포함할 수 있습니다. "serving_default"는 저장된 케라스 모델이 있는 추론 기능을 위한 기본 키입니다. 이 기능을 이용하여 추론합니다:


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)

  # 분산방식으로 기능 호출하기
  for batch in dist_predict_dataset:
    another_strategy.run(inference_func,args=(batch,))

복원된 기능을 호출하는 것은 단지 저장된 모델로의 정방향 패쓰입니다(예상하기에). 만약 계속해서 불러온 기능을 훈련시키고 싶다면 어떻게 하실건가요? 불러온 기능을 더 큰 모델에 내장시킬 건가요? 일반적인 방법은 이 불러온 객체를 케라스 층에 싸서(wrap) 달성하는 것입니다. 다행히도, 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')
  # 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=['accuracy'])
  model.fit(train_dataset, epochs=2)

볼 수 있듯이, hub.KerasLayertf.saved_model.load()로부터 불려온 결과를 또 다른 모델을 만드는데 사용될 수 있는 케라스 층으로 포장(wrap)합니다. 이것은 학습에 매우 유용합니다.

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

저장에 관해서, 케라스 모델을 사용하는 경우, 케라스의 model.save() API를 사용하는 것을 권장합니다. 저장하려는 모델이 케라스 모델이 아닌 경우, 더 낮은 단계의 API를 선택해야 합니다.

모델을 불러옴에 있어서, 어떤 API를 사용하느냐는 로딩 API에서 얻고자 하는 내용에 따라 결정됩니다. 케라스 모델을 가져올 수 없으면(또는 가져오고 싶지 않다면), tf.saved_model.load()를 사용합니다. 그 외의 경우에는, tf.keras.models.load_model()을 사용합니다. 케라스 모델을 저장한 경우에만 케라스 모델을 반환 받을 수 있다는 점을 유의하세요.

API들을 목적에 따라 혼합하고 짜 맞추는 것이 가능합니다. 케라스 모델을 model.save와 함께 저장할 수 있고, 저수준 API인, tf.saved_model.load로 케라스가 아닌 모델을 불러올 수 있습니다.


In [ ]:
model = get_model()

# 케라스의 save() API를 사용하여 모델 저장하기
model.save(keras_model_path) 

another_strategy = tf.distribute.MirroredStrategy()
# 저수준 API를 사용하여 모델 불러오기
with another_strategy.scope():
  loaded = tf.saved_model.load(keras_model_path)

주의사항

특별한 경우는 잘 정의되지 않은 입력을 갖는 케라스 모델을 갖고 있는 경우입니다. 예를 들어, 순차 모델은 입력 형태(Sequential([Dense(3), ...]) 없이 만들 수 있습니다. 하위 분류된 모델들 또한 초기화 후에 잘 정의된 입력을 갖고 있지 않습니다. 이 경우 모델을 저장하고 불러올 시 저수준 API를 사용해야 하며, 그렇지 않으면 오류가 발생할 수 있습니다.

모델이 잘 정의된 입력을 갖는지 확인하려면, model.inputsNone인지 확인합니다. None이 아니라면 잘 정의된 입력입니다. 입력 형태들은 모델이 .fit, .evaluate, .predict에서 쓰이거나 모델을 호출 (model(inputs)) 할 때 자동으로 정의됩니다.

예시를 살펴봅시다:


In [ ]:
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)  # 오류! 
tf.saved_model.save(my_model, saved_model_path)