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.

Guardar y cargar un modelo usando una estrategia de distribución

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

Visión general

En este tutorial se muestra el modo de guardar y cargar modelos en formato SavedModel con tf.distribute.Strategy durante o después del entrenamiento. Hay dos tipos de API para guardar y cargar un modelo Keras: de alto nivel (tf.keras.Model.save y tf.keras.models.load_model) y de bajo nivel (tf.saved_model.save y tf.saved_model.load).

Para aprender sobre SavedModel y la serialización en general, lea la guía de modelo guardado, y la guía de serialización de modelo Keras. Empecemos con un ejemplo simple.

Advertencia: Los modelos TensorFlow son código y debe tener cuidado con el código que no sea de confianza. Aprenda más en Usar TensorFlow de forma segura.

Importar dependencias:

import tensorflow_datasets as tfds

import tensorflow as tf

Cargue y prepare los datos con conjuntos de datos TensorFlow y tf.data, y cree el modelo utilizando 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

Entrene el modelo con tf.keras.Model.fit:

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

Guardar y cargar el modelo

Ahora que tiene un modelo sencillo con el que trabajar, exploremos las API de guardado/carga. Hay dos tipos de API disponibles:

• Alto nivel (Keras): Model.save y tf.keras.models.load_model (formato de archivo zip .keras)

• Bajo nivel: tf.saved_model.save y tf.saved_model.load (formato TF SavedModel)

La API Keras

He aquí un ejemplo de cómo guardar y cargar un modelo con la API de Keras:

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

Restaure el modelo sin tf.distribute.Strategy:

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

Una vez restaurado el modelo, puede seguir entrenando sobre él, incluso sin necesidad de volver a llamar a Model.compile, puesto que ya estaba compilado antes de guardarlo. El modelo se guarda en formato de archivo zip de Keras, marcado con la extensión .keras. Para más información, consulte la guía sobre guardado de Keras.

Ahora, restaure el modelo y entrénelo usando una 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 muestra la salida Model.fit, la carga funciona como se esperaba con tf.distribute.Strategy. La estrategia usada aquí no tiene por qué ser la misma usada antes de guardar.

La API tf.saved_model

Guardar el modelo con la API de nivel inferior es parecido a hacerlo con la API de Keras:

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

Se puede cargar con tf.saved_model.load. Sin embargo, dado que se trata de una API de nivel más bajo (y, por tanto, tiene una gama más amplia de casos de uso), no devuelve un modelo Keras. En su lugar, devuelve un objeto que contiene funciones que pueden usarse para hacer inferencias. Por ejemplo:

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

El objeto cargado puede contener varias funciones, cada una asociada a una clave. La clave "serving_default" es la clave predeterminada para la función de inferencia con un modelo Keras guardado. Para realizar la inferencia con esta función:

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

También puede cargar y realizar inferencias de forma distribuida:

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

Llamar a la función restaurada es sólo un pase hacia adelante sobre el modelo guardado (tf.keras.Model.predict). ¿Y si desea seguir entrenando la función cargada? ¿O qué pasa si necesita incrustar la función cargada en un modelo más grande? Una práctica común es empaquetar este objeto cargado en una capa Keras para conseguirlo. Por suerte, TF Hub dispone de hub.KerasLayer con este fin, como se muestra aquí:

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)

En el ejemplo anterior, hub.KerasLayer de Tensorflow Hub envuelve el resultado cargado de nuevo desde tf.saved_model.load en una capa Keras que se usa para construir otro modelo. Esto es muy útil para el aprendizaje por transferencia.

¿Qué API debiera usar?

Para guardar, si está trabajando con un modelo Keras, use el Keras Model.save a menos que necesite el control adicional que permite la API de bajo nivel. Si lo que está guardando no es un modelo Keras, entonces la API de bajo nivel, tf.saved_model.save, es su única opción.

Para cargar, su elección de API depende de lo que quiera recibir de la API de carga de modelos. Si no puede (o no quiere) recibir un modelo Keras, entonces use tf.saved_model.load. De lo contrario, use tf.keras.models.load_model. Tenga en cuenta que sólo puede recuperar un modelo Keras si ha guardado un modelo Keras.

Es posible combinar las API. Puede guardar un modelo Keras con Model.save, y cargar un modelo que no sea Keras con la API de bajo nivel, 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)

Guardar/cargar desde un dispositivo local

Al guardar y cargar desde un dispositivo de E/S local mientras se entrena en dispositivos remotos (por ejemplo, al usar una TPU Cloud), debe usar la opción experimental_io_device en tf.saved_model.SaveOptions y tf.saved_model.LoadOptions para fijar el dispositivo de E/S en localhost. Por ejemplo:

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)

Precauciones

Un caso especial es cuando se crean modelos Keras de determinadas formas y luego se guardan antes de entrenar. Por ejemplo:

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)

Un SavedModel guarda los objetos tf.types.experimental.ConcreteFunction generados al rastrear una tf.function (consulte ¿Cuándo se rastrea una Función? en la guía Introducción a los grafos y tf.function para obtener más información). Si obtiene un ValueError como este, es porque Model.save no fue capaz de encontrar o crear una ConcreteFunction rastreada.

Precaución: No debería guardar un modelo sin al menos una ConcreteFunction, ya que, de lo contrario, la API de bajo nivel generará un SavedModel sin firmas ConcreteFunction (aprenda más sobre el formato SavedModel). Por ejemplo:

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

Normalmente, la pasada hacia delante del modelo (el método call) se rastrea automáticamente cuando se llama al modelo por primera vez, a menudo a través del método Model.fit de Keras. Una ConcreteFunction también puede ser generada por las API Sequential y Functional de Keras, si se configura la forma de entrada, por ejemplo, haciendo la primer capa ya sea un tf.keras.layers.InputLayer u otro tipo de capa, y pasándole el argumento de palabra clave input_shape.

Para verificar si su modelo tiene alguna ConcreteFunction rastreada, compruebe si Model.save_spec es None:

print(my_model.save_spec() is None)

Usemos tf.keras.Model.fit para entrenar el modelo, y notemos que el save_spec queda definido y el guardado del 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)