Asista al Simposio Mujeres en ML el 7 de diciembre Regístrese ahora

Modelos guardados de TF Hub en TensorFlow 2

Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.

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

Esta página explica cómo reutilizar modelos guardados de TF2 en un programa TensorFlow 2 con la API hub.load() de bajo nivel y su contenedor hub.KerasLayer . (Normalmente, hub.KerasLayer se combina con otros tf.keras.layers para construir un modelo Keras o el model_fn de un TF2 Estimator). Estas API también pueden cargar los modelos heredados en formato TF1 Hub, dentro de los 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.

Uso de 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 de capa de Keras. La biblioteca tensorflow_hub proporciona la clase hub.KerasLayer que se inicializa con la URL (o la ruta del sistema de archivos) de un modelo guardado y luego proporciona el cálculo del modelo guardado, incluidos sus pesos preentrenados.

Aquí hay 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 en la forma habitual de Keras:

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

La colaboración 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 ajustes 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.

Uso de un modelo guardado en un estimador

Los usuarios de la API Estimator de TensorFlow para entrenamiento distribuido pueden usar modelos guardados 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 modelos guardados

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 modelos guardados descargados y sin comprimir. Para obtener más información, consulte Almacenamiento en caché .

Usando un modelo guardado en TensorFlow de bajo nivel

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

Tema avanzado: qué esperar del modelo guardado después de cargar

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

  • Las firmas de servicio del modelo guardado (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 la entrada y salida respectivas nombres y sujeto a las restricciones de forma y tipo de firma.

  • Los métodos @tf.function del objeto guardado (si los hay) se restauran como objetos tf.function que pueden ser llamados por todas las combinaciones de argumentos Tensor y no Tensor para los que se había rastreado la función tf.antes de guardar. En particular, si hay un método obj.__call__ con trazas adecuadas, 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 puede implementar SavedModels. La interfaz de modelos guardados reutilizables para obj establece convenciones de modo que el código del cliente, incluidos los adaptadores como hub.KerasLayer , sepa cómo usar el modelo guardado.

Es posible que algunos modelos guardados 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 modelo guardado se vuelven a cargar como entrenables y tf.GradientTape las observará de manera 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 aconseja volver a entrenar solo un subconjunto de las variables originalmente entrenables.

Creación de modelos guardados para TF Hub

Visió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 usar 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.

Salvando de Keras

A partir de TensorFlow 2, tf.keras.Model.save() y tf.keras.models.save_model() por defecto el formato de modelo guardado (no HDF5). Los modelos guardados resultantes que 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 de Keras completo, simplemente guárdelo con include_optimizer=False .

Para compartir una parte de un modelo de Keras, haga que la parte sea un modelo en sí mismo y luego guárdelo. Puede 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 recorte la pieza para compartirla después (si se alinea con las capas de su 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 usa el enfoque anterior 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 ).

Guardar desde TensorFlow de bajo nivel

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

Si desea proporcionar algo más que una firma de publicación, debe implementar la interfaz Reusable SavedModel . 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 modelo guardado importado junto con las del modelo que lo rodea se denomina ajuste fino del modelo guardado. 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, aumentar el riesgo de sobreajuste y requerir un aumento del conjunto de datos, especialmente para las CNN). Aconsejamos a los consumidores de SavedModel que analicen el ajuste fino 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 tokenizar la entrada de texto y asignar tokens a sus entradas correspondientes en una matriz de incrustación.

Para consumidores de modelos guardados

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 el modelo guardado al modelo de Keras y ejecuta el cálculo del modelo guardado en modo de entrenamiento (piense en el abandono, etc.).

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

Reexportación del resultado del ajuste fino

Es posible que los usuarios avanzados deseen guardar los resultados del ajuste fino en un modelo guardado que se pueda usar en lugar del que se cargó originalmente. Esto se puede hacer con un 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 modelos guardados

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

Guardar desde un modelo de Keras debería hacer que todos los mecanismos de ajuste fino funcionen (ahorrar pérdidas de regularización de peso, declarar variables entrenables, rastrear __call__ tanto para training=True como para 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 de destino esperados y tamaños de lotes. (Al momento de escribir este artículo, guardar de 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 de su modelo guardado en consecuencia.