Modelos guardados de TF Hub en TensorFlow 2

El formato de SavedModel TensorFlow 2 es la forma recomendada para compartir modelos pre-formados y piezas modelo sobre TensorFlow concentradores. Se reemplaza el viejo formato de TF1 Hub y viene con un nuevo conjunto de APIs.

Esta página explica cómo reutilizar TF2 SavedModels en un programa TensorFlow 2 con el bajo nivel de hub.load() de la API y su hub.KerasLayer envoltura. (Normalmente, hub.KerasLayer se combina con otros tf.keras.layers para construir un modelo Keras o la model_fn de un TF2 Estimador.) 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 API de alto nivel de TensorFlow para la construcción de modelos de aprendizaje profundo mediante la composición de los objetos de capa Keras. El tensorflow_hub biblioteca proporciona la clase hub.KerasLayer que consigue inicializa con el URL (o ruta de sistema de archivos) de un SavedModel y luego proporciona el cálculo de la SavedModel, incluyendo sus pesos pre-formado.

A continuación, se muestra un ejemplo del uso de una inserció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 de Keras:

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

El colab clasificación de texto es un ejemplo completo de cómo entrenar y evaluar a un clasificador tal.

Los pesos de modelo en un hub.KerasLayer se establecen para no entrenable de forma predeterminada. Consulte la sección de 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.

Usar un modelo guardado en un estimador

Los usuarios de TensorFlow del Estimador de API para la formación distribuida pueden utilizar 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 modelos guardados

El uso de un SavedModel de TensorFlow Hub (u otros servidores HTTPS que implementan sus recibimiento del protocolo) y descomprime las descargas al sistema de archivos local si no está ya 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 descomprimidos. Para más detalles, véase el almacenamiento en caché .

Usar un modelo guardado en TensorFlow de bajo nivel

La función hub.load(handle) descargas y descomprime un SavedModel (a menos que handle es ya una trayectoria de sistema de archivos) y luego devuelve el resultado de cargarlo con función integrada de TensorFlow tf.saved_model.load() . Por lo tanto, hub.load() puede manejar cualquier SavedModel válida (a diferencia de su predecesor hub.Module para TF1).

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

En función de los contenidos de la SavedModel, el resultado de obj = hub.load(...) se puede invocar de diversas maneras (como se explica en mayor detalle en la mayor de TensorFlow Guía SavedModel :

  • Las firmas de porción de la SavedModel (si los hay) se representan como un diccionario de funciones concretas y pueden ser llamados como tensors_out = obj.signatures["serving_default"](**tensors_in) , con los diccionarios de los tensores tecleados por la entrada y la salida respectiva nombres y sujeto a las restricciones de forma y tipo de la firma.

  • Los @tf.function métodos -decorated del objeto guardado (si los hay) se restauran como objetos tf.function que pueden ser llamados por todas las combinaciones de Tensor y no Tensor argumentos para el cual el tf.function había sido trazadas antes de guardar. En particular, si hay un obj.__call__ método con trazas adecuados, obj en sí puede ser llamado como una función de Python. Un ejemplo simple podría parecerse a output_tensor = obj(input_tensor, training=False) .

Esto deja una enorme libertad en las interfaces que pueden implementar SavedModels. La interfaz reutilizables SavedModels para obj establece convenciones de tal manera que el código del cliente, incluyendo los adaptadores como hub.KerasLayer , sabe cómo utilizar el SavedModel.

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 servicio.

Las variables entrenables en un SavedModel se vuelven a cargar como entrenable, y tf.GradientTape serán verlas de forma predeterminada. Consulte la sección sobre ajustes a continuación para ver algunas advertencias y, para empezar, considere evitar esto. Incluso si usted quiere poner a punto, es posible que desee ver si obj.trainable_variables aconseja a re-tren sólo un subconjunto de las variables originalmente entrenables.

Creación de 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.

Ahorrando desde Keras

A partir de TensorFlow 2, tf.keras.Model.save() y tf.keras.models.save_model() por defecto para el formato SavedModel (no HDF5). Los SavedModels resultantes que se pueden utilizar 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 completa, simplemente guardarlo con include_optimizer=False .

Para compartir una pieza de un modelo Keras, convierta la pieza en 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 recorta la pieza para compartir después del hecho (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 Modelos en GitHub utiliza el primer enfoque para el BERT (ver PNL / herramientas / export_tfhub_lib.py , tenga en cuenta la división entre core_model para la exportación y la pretrainer para restaurar el punto de control) y el último enfoque para ResNet (ver la visión / image_classification / tfhub_export. PY ).

Guardar desde TensorFlow de bajo nivel

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

Si desea proporcionar algo más que una firma de servir, se debe implementar la interfaz reutilizable 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

El entrenamiento de las variables ya formados, de una SavedModel importados junto con los del modelo alrededor de ella se llama ajuste fino del SavedModel. Esto puede resultar en 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 consideren 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 de forma rígida, como la creación de tokens de entrada de texto y la asignación de tokens a sus entradas correspondientes en una matriz de incrustación.

Para los consumidores de SavedModel

La creación de 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 el modo de entrenamiento (piense en la deserción, etc.).

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

Reexportar el resultado del ajuste fino

Los usuarios avanzados pueden querer guardar los resultados del ajuste fino de nuevo en un modelo guardado que se puede 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 en el futuro si sus consumidores deberían ajustarlo y en qué forma, y ​​brinde orientación en la documentación.

El ahorro de un modelo Keras debe hacer toda la mecánica de trabajo de puesta a punto (ahorro de las pérdidas de peso de regularización, la declaración de variables entrenables, trazando __call__ tanto para training=True y 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 usa abandono, normalización de 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 lote. (Al momento de escribir este artículo, ahorrar con Keras no facilita que los consumidores los ajusten).

Regularizers peso de las capas individuales se guardan (con sus coeficientes de resistencia regularización), pero regularización de peso dentro del optimizador (como tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) se pierde. Avise a los consumidores de su modelo guardado en consecuencia.