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.

Salve e carregue um modelo usando uma estratégia de distribuição

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

Este tutorial demonstra como você pode salvar e carregar modelos no formato SavedModel com tf.distribute.Strategy durante ou após o treinamento. Existem dois tipos de API para salvar e carregar um modelo do Keras: um de alto nível (tf.keras.Model.save e tf.keras.models.load_model) e outro de baixo nível (tf.saved_model.save e tf.saved_model.load).

Para saber mais sobre SavedModel e serialização de forma geral, leia o guia SavedModel e o guia Serialização do modelo do Keras. Vamos começar com um exemplo simples.

Atenção: os modelos do TensorFlow são códigos, e é importante ter cuidado com código não confiável. Saiba mais em Como usar o TensorFlow com segurança.

Importe as dependências:

import tensorflow_datasets as tfds

import tensorflow as tf

Carregue e prepare os dados com os datasets do TensorFlow e tf.data, e crie o modelo usando 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

Treine o modelo com tf.keras.Model.fit:

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

Como salvar e carregar o modelo

Agora que você já tem um modelo simples, vamos conferir as APIs de salvamento/carregamento. Dois tipos de API estão disponíveis:

• Alto nível (Keras): Model.save e tf.keras.models.load_model (formato de arquivo zip .keras)

• Baixo nível: tf.saved_model.save e tf.saved_model.load (formato SavedModel do TF)

API do Keras

Veja um exemplo de como salvar e carregar um modelo com a API do Keras:

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

Restaure o modelo sem tf.distribute.Strategy:

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

Após restaurar o modelo, você pode continuar fazendo o treinamento, mesmo sem precisar fazer uma chamada a Model.compile novamente, pois ele já foi compilado antes do salvamento. O modelo é salvo no formato de arquivo zip do Keras, marcado pela extensão .keras. Confira mais informações no guia sobre salvamento do Keras.

Agora, restaure o modelo e faça o treinamento usando 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)

Como a saída de Model.fit mostra, o carregamento funciona conforme o esperado com tf.distribute.Strategy. A estratégia usada aqui não precisa ser a mesma usada antes do salvamento.

API tf.saved_model

Salvar o modelo com a API de baixo nível é similar ao salvamento com a API do Keras:

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

O carregamento pode ser feito com tf.saved_model.load. Entretanto, como é uma API de baixo nível (e, portanto, tem uma gama maior de casos de uso), ela não retorna um modelo do Keras. Em vez disso, retorna um objeto que contém funções que podem ser usadas para fazer inferência. Por exemplo:

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

O objeto carregado pode conter diversas funções, cada uma associada a uma chave. A chave "serving_default" é a chave padrão para a função de inferência com um modelo do Keras salvo. Para fazer inferência com essa função:

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

Você também pode carregar e fazer inferência de uma forma distribuída:

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

A chamada à função restaurada é apenas um passo para frente do modelo salvo (tf.keras.Model.predict). E se você quiser continuar treinando a função carregada? Ou se você precisar embutir uma função carregada em um modelo maior? Uma prática comum para fazer isso é encapsular esse objeto carregado em uma camada do Keras. Felizmente, o TF Hub tem hub.KerasLayer para essa finalidade. Veja:

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)

No exemplo acima, o hub.KerasLayer do TensorFlow Hub encapsula o resultado carregado a partir de tf.saved_model.load em uma camada do Keras, que é usada para criar outro modelo. Isso é muito útil para aprendizado por transferência.

Que API devo usar?

Para salvar, se você estiver trabalhando com um modelo do Keras, use a API Model.save do Keras, a menos que você precise de controle adicional permitido pela API de nível baixo. Se o que você estiver salvando não for um modelo do Keras, então a API de baixo nível, tf.saved_model.save, é sua única opção.

Para carregar, a escolha da API depende do que você deseja obter com a API de carregamento do modelo. Se você não puder (ou não quiser) obter um modelo do Keras, use tf.saved_model.load. Caso contrário, use tf.keras.models.load_model. Você pode obter um modelo do Keras somente se tiver salvo um.

É possível usar as duas APIs ao mesmo tempo. Você pode salvar um modelo do Keras com Model.save e carregar um modelo que não é do Keras com a API de baixo nível, tf.saved_model.load.

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)

Como salvar/carregar usando um dispositivo local

Ao salvar e carregar usando um dispositivos de E/S local ao fazer o treinamento em dispositivos remotos (por exemplo, ao usar TPU em nuvem), você precisa usar a opção experimental_io_device em tf.saved_model.SaveOptions e tf.saved_model.LoadOptions para definir o dispositivo de E/S como localhost. Por exemplo:

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)

Ressalvas

Um caso especial é quando você cria modelos do Keras de determinadas formas e depois salva-os antes do treinamento. Por exemplo:

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 salva os objetos tf.types.experimental.ConcreteFunction gerados quando você faz o trace de uma tf.function (confira Quando uma função está fazendo trace? no guia Introdução a grafos e tf.function para saber mais). Se você obtiver um ValueError, é porque Model.save não conseguiu encontrar ou criar uma ConcreteFunction que sofreu trace.

Atenção: você não deve salvar um modelo sem pelo menos uma ConcreteFunction, já que, caso o faça, a API de baixo nível gerará um SavedModel sem assinaturas de ConcreteFunction (saiba mais sobre o formato SavedModel). Por exemplo:

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

Geralmente, o passo para frente do modelo (o método call) sofrerá trace automaticamente quando o modelo for chamado pela primeira vez, geralmente pelo método Model.fit do Keras. Uma ConcreteFunction também pode ser gerada pelas APIs Sequential e Functional do Keras se você definir o formato da entrada – por exemplo, usando como primeira camada um tf.keras.layers.InputLayer ou outro tipo de camada – e passando a ela o argumento palavra-chave input_shape.

Para verificar se seu modelo tem alguma ConcreteFunction que sofreu trace, confira se Model.save_spec é igual a None (nenhuma):

print(my_model.save_spec() is None)

Vamos usar tf.keras.Model.fit para treinar o modelo. Observe que save_spec é definido, e o salvamento do modelo funcionará:

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)