Modelos guardados de TF Hub en TensorFlow 2

El formato SavedModel de TensorFlow 2 es la forma recomendada de compartir modelos y piezas de modelo previamente entrenados en TensorFlow Hub. Reemplaza el antiguo formato TF1 Hub y viene con un nuevo conjunto de API.

Esta página explica cómo reutilizar TF2 SavedModels en un programa TensorFlow 2 con la API hub.load() de bajo nivel y su contenedor hub.KerasLayer . (Por lo general, hub.KerasLayer se combina con otros tf.keras.layers para crear un modelo Keras o el model_fn de un Estimador TF2). Estas API también pueden cargar los modelos heredados en formato TF1 Hub, dentro de ciertos límites; consulte la guía de compatibilidad .

Los usuarios de TensorFlow 1 pueden actualizar a TF 1.15 y luego usar las mismas API. Las versiones anteriores de TF1 no funcionan.

Usando modelos guardados de TF Hub

Usando un modelo guardado en Keras

Keras es la API de alto nivel de TensorFlow para crear modelos de aprendizaje profundo mediante la composición de objetos Keras Layer. La biblioteca tensorflow_hub proporciona la clase hub.KerasLayer que se inicializa con la URL (o ruta del sistema de archivos) de un SavedModel y luego proporciona el cálculo del SavedModel, incluidos sus pesos previamente entrenados.

A continuación se muestra un ejemplo del uso de una incrustación de texto previamente entrenada:

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://tfhub.dev/google/nnlm-en-dim128/2"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

A partir de esto, se puede construir un clasificador de texto de la forma habitual en Keras:

model = tf.keras.Sequential([
    embed,
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(1, activation="sigmoid"),
])

La colab de clasificación de texto es un ejemplo completo de cómo entrenar y evaluar dicho clasificador.

Los pesos del modelo en un hub.KerasLayer están configurados como no entrenables de forma predeterminada. Consulte la sección sobre ajuste fino a continuación para saber cómo cambiar eso. Los pesos se comparten entre todas las aplicaciones del mismo objeto de capa, como es habitual en Keras.

Usando un modelo guardado en un estimador

Los usuarios de la API Estimator de TensorFlow para entrenamiento distribuido pueden usar SavedModels de TF Hub escribiendo su model_fn en términos de hub.KerasLayer entre otros tf.keras.layers .

Detrás de escena: descarga y almacenamiento en caché de SavedModel

El uso de un modelo guardado de TensorFlow Hub (u otros servidores HTTPS que implementen su protocolo de alojamiento ) lo descarga y lo descomprime en el sistema de archivos local si aún no está presente. La variable de entorno TFHUB_CACHE_DIR se puede configurar para anular la ubicación temporal predeterminada para almacenar en caché los SavedModels descargados y sin comprimir. Para obtener más información, consulte Almacenamiento en caché .

Usando un modelo guardado en TensorFlow de bajo nivel

Manijas modelo

SavedModels se puede cargar desde un handle específico, donde el handle es una ruta del sistema de archivos, una URL válida del modelo TFhub.dev (por ejemplo, "https://tfhub.dev/..."). Las URL de Kaggle Models reflejan los manejadores de TFhub.dev de acuerdo con nuestros Términos y la licencia asociada con los activos del modelo, por ejemplo, "https://www.kaggle.com/...". Los identificadores de los modelos Kaggle son equivalentes a su identificador TFhub.dev correspondiente.

La función hub.load(handle) descarga y descomprime un SavedModel (a menos que handle ya sea una ruta del sistema de archivos) y luego devuelve el resultado de cargarlo con la función incorporada de TensorFlow tf.saved_model.load() . Por lo tanto, hub.load() puede manejar cualquier SavedModel válido (a diferencia de su predecesor hub.Module para TF1).

Tema avanzado: qué esperar del SavedModel después de cargarlo

Dependiendo del contenido de SavedModel, el resultado de obj = hub.load(...) se puede invocar de varias maneras (como se explica con mucho más detalle en la Guía SavedModel de TensorFlow:

  • Las firmas de servicio de SavedModel (si las hay) se representan como un diccionario de funciones concretas y se pueden llamar como tensors_out = obj.signatures["serving_default"](**tensors_in) , con diccionarios de tensores codificados por las respectivas entradas y salidas. nombres y sujeto a las restricciones de forma y tipo de firma.

  • Los métodos decorados con @tf.function del objeto guardado (si los hay) se restauran como objetos tf.function que pueden ser invocados por todas las combinaciones de argumentos tensoriales y no tensoriales para los cuales se ha rastreado la función tf.function antes de guardar. En particular, si hay un método obj.__call__ con seguimientos adecuados, el propio obj se puede llamar como una función de Python. Un ejemplo simple podría verse como output_tensor = obj(input_tensor, training=False) .

Esto deja una enorme libertad en las interfaces que SavedModels puede implementar. La interfaz reutilizable SavedModels para obj establece convenciones tales que el código del cliente, incluidos los adaptadores como hub.KerasLayer , sabe cómo utilizar SavedModel.

Es posible que algunos SavedModels no sigan esa convención, especialmente los modelos completos que no están destinados a ser reutilizados en modelos más grandes y solo proporcionan firmas de publicación.

Las variables entrenables en un SavedModel se recargan como entrenables y tf.GradientTape las observará de forma predeterminada. Consulte la sección sobre ajustes a continuación para conocer algunas advertencias y considere evitar esto para empezar. Incluso si desea realizar un ajuste fino, es posible que desee ver si obj.trainable_variables recomienda volver a entrenar solo un subconjunto de las variables originalmente entrenables.

Creando modelos guardados para TF Hub

Descripción general

SavedModel es el formato de serialización estándar de TensorFlow para modelos entrenados o piezas de modelos. Almacena los pesos entrenados del modelo junto con las operaciones exactas de TensorFlow para realizar su cálculo. Se puede utilizar independientemente del código que lo creó. En particular, se puede reutilizar en diferentes API de creación de modelos de alto nivel como Keras, porque las operaciones de TensorFlow son su lenguaje básico común.

Ahorro de Keras

A partir de TensorFlow 2, tf.keras.Model.save() y tf.keras.models.save_model() tienen de forma predeterminada el formato SavedModel (no HDF5). Los SavedModels resultantes se pueden usar con hub.load() , hub.KerasLayer y adaptadores similares para otras API de alto nivel a medida que estén disponibles.

Para compartir un modelo Keras completo, simplemente guárdelo con include_optimizer=False .

Para compartir una parte de un modelo Keras, convierta la pieza en un modelo en sí misma y luego guárdela. Puedes diseñar el código así desde el principio...

piece_to_share = tf.keras.Model(...)
full_model = tf.keras.Sequential([piece_to_share, ...])
full_model.fit(...)
piece_to_share.save(...)

...o recorta la pieza para compartirla después (si se alinea con las capas de tu modelo completo):

full_model = tf.keras.Model(...)
sharing_input = full_model.get_layer(...).get_output_at(0)
sharing_output = full_model.get_layer(...).get_output_at(0)
piece_to_share = tf.keras.Model(sharing_input, sharing_output)
piece_to_share.save(..., include_optimizer=False)

TensorFlow Models en GitHub utiliza el primer enfoque para BERT (consulte nlp/tools/export_tfhub_lib.py , tenga en cuenta la división entre core_model para exportar y el pretrainer para restaurar el punto de control) y el último enfoque para ResNet (consulte Legacy/image_classification/tfhub_export.py ).

Ahorro desde TensorFlow de bajo nivel

Esto requiere una buena familiaridad con la Guía SavedModel de TensorFlow.

Si desea proporcionar algo más que una firma de entrega, debe implementar la interfaz de modelo guardado reutilizable . Conceptualmente, esto parece

class MyMulModel(tf.train.Checkpoint):
  def __init__(self, v_init):
    super().__init__()
    self.v = tf.Variable(v_init)
    self.variables = [self.v]
    self.trainable_variables = [self.v]
    self.regularization_losses = [
        tf.function(input_signature=[])(lambda: 0.001 * self.v**2),
    ]

  @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
  def __call__(self, inputs):
    return tf.multiply(inputs, self.v)

tf.saved_model.save(MyMulModel(2.0), "/tmp/my_mul")

layer = hub.KerasLayer("/tmp/my_mul")
print(layer([10., 20.]))  # [20., 40.]
layer.trainable = True
print(layer.trainable_weights)  # [2.]
print(layer.losses)  # 0.004

Sintonia FINA

Entrenar las variables ya entrenadas de un SavedModel importado junto con las del modelo que lo rodea se denomina ajuste fino de SavedModel. Esto puede dar como resultado una mejor calidad, pero a menudo hace que el entrenamiento sea más exigente (puede llevar más tiempo, depender más del optimizador y sus hiperparámetros, aumenta el riesgo de sobreajuste y requiere un aumento del conjunto de datos, especialmente para las CNN). Aconsejamos a los consumidores de SavedModel que consideren realizar ajustes solo después de haber establecido un buen régimen de capacitación y solo si el editor de SavedModel lo recomienda.

El ajuste fino cambia los parámetros del modelo "continuo" que se entrenan. No cambia las transformaciones codificadas, como la tokenización de la entrada de texto y la asignación de tokens a sus entradas correspondientes en una matriz de incrustación.

Para consumidores de SavedModel

Creando un hub.KerasLayer como

layer = hub.KerasLayer(..., trainable=True)

permite el ajuste fino del modelo guardado cargado por la capa. Agrega los pesos entrenables y los regularizadores de peso declarados en SavedModel al modelo Keras y ejecuta el cálculo de SavedModel en modo de entrenamiento (piense en el abandono, etc.).

La colab de clasificación de imágenes contiene un ejemplo de un extremo a otro con ajuste fino opcional.

Reexportar el resultado del ajuste fino

Es posible que los usuarios avanzados quieran guardar los resultados del ajuste fino en un SavedModel que pueda usarse en lugar del cargado originalmente. Esto se puede hacer con código como

loaded_obj = hub.load("https://tfhub.dev/...")
hub_layer = hub.KerasLayer(loaded_obj, trainable=True, ...)

model = keras.Sequential([..., hub_layer, ...])
model.compile(...)
model.fit(...)

export_module_dir = os.path.join(os.getcwd(), "finetuned_model_export")
tf.saved_model.save(loaded_obj, export_module_dir)

Para creadores de SavedModel

Al crear un modelo guardado para compartir en TensorFlow Hub, piense con anticipación si sus consumidores deberían ajustarlo y cómo, y proporcione orientación en la documentación.

Guardar desde un modelo Keras debería hacer que todos los mecanismos de ajuste funcionen (ahorrar pérdidas de regularización de peso, declarar variables entrenables, rastrear __call__ tanto para training=True como training=False , etc.)

Elija una interfaz de modelo que funcione bien con el flujo de gradiente, por ejemplo, logits de salida en lugar de probabilidades softmax o predicciones top-k.

Si el modelo utiliza abandono, normalización por lotes o técnicas de entrenamiento similares que involucran hiperparámetros, configúrelos en valores que tengan sentido en muchos problemas objetivo y tamaños de lote esperados. (En el momento de escribir este artículo, ahorrar en Keras no facilita que los consumidores los ajusten).

Los regularizadores de peso en capas individuales se guardan (con sus coeficientes de fuerza de regularización), pero la regularización de peso desde dentro del optimizador (como tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) se pierde. Informe a los consumidores sobre su SavedModel en consecuencia.