Modèles enregistrés de TF Hub dans TensorFlow 2

Le format SavedModel de TensorFlow 2 est la méthode recommandée pour partager des modèles et des éléments de modèle pré-entraînés sur TensorFlow Hub. Il remplace l'ancien format TF1 Hub et est livré avec un nouvel ensemble d'API.

Cette page explique comment réutiliser TF2 SavedModels dans un programme TensorFlow 2 avec l'API bas niveau hub.load() et son wrapper hub.KerasLayer . (Typiquement, hub.KerasLayer est combiné avec d'autres tf.keras.layers pour construire un modèle Keras ou le model_fn d'un estimateur TF2.) Ces API peuvent également charger les anciens modèles au format TF1 Hub, dans certaines limites, voir le guide de compatibilité .

Les utilisateurs de TensorFlow 1 peuvent mettre à jour vers TF 1.15, puis utiliser les mêmes API. Les anciennes versions de TF1 ne fonctionnent pas.

Utilisation de SavedModels à partir de TF Hub

Utilisation d'un modèle enregistré dans Keras

Keras est l'API de haut niveau de TensorFlow pour créer des modèles d'apprentissage en profondeur en composant des objets Keras Layer. La bibliothèque tensorflow_hub fournit la classe hub.KerasLayer qui est initialisée avec l'URL (ou le chemin du système de fichiers) d'un SavedModel, puis fournit le calcul à partir du SavedModel, y compris ses poids pré-formés.

Voici un exemple d'utilisation d'une intégration de texte pré-entraînée :

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)

À partir de là, un classificateur de texte peut être construit de la manière habituelle de Keras :

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

Le Colab de classification de texte est un exemple complet de formation et d'évaluation d'un tel classifieur.

Les pondérations du modèle dans un hub.KerasLayer sont définies sur non entraînables par défaut. Voir la section sur le réglage ci-dessous pour savoir comment changer cela. Les poids sont partagés entre toutes les applications du même objet de calque, comme d'habitude dans Keras.

Utilisation d'un modèle enregistré dans un estimateur

Les utilisateurs de l'API Estimator de TensorFlow pour la formation distribuée peuvent utiliser SavedModels de TF Hub en écrivant leur model_fn en termes de hub.KerasLayer parmi d'autres tf.keras.layers .

Dans les coulisses : téléchargement et mise en cache de SavedModel

L'utilisation d'un SavedModel de TensorFlow Hub (ou d'autres serveurs HTTPS qui implémentent son protocole d'hébergement ) le télécharge et le décompresse dans le système de fichiers local s'il n'est pas déjà présent. La variable d'environnement TFHUB_CACHE_DIR peut être définie pour remplacer l'emplacement temporaire par défaut pour la mise en cache des SavedModels téléchargés et non compressés. Pour plus de détails, consultez Mise en cache .

Utiliser un SavedModel dans TensorFlow de bas niveau

Modèle Poignées

Les SavedModels peuvent être chargés à partir d'un handle spécifié, où le handle est un chemin de système de fichiers, une URL de modèle TFhub.dev valide (par exemple "https://tfhub.dev/ ..."). Les URL des modèles de Kaggle reflètent les poignées de TFhub.dev conformément à nos Conditions et à la licence associée aux actifs du modèle, par exemple "https://www.kaggle.com/...". Les handles des modèles Kaggle sont équivalents à leur handle TFhub.dev correspondant.

La fonction hub.load(handle) télécharge et décompresse un SavedModel (sauf si handle est déjà un chemin de système de fichiers), puis renvoie le résultat du chargement avec la fonction intégrée de TensorFlow tf.saved_model.load() . Par conséquent, hub.load() peut gérer n'importe quel SavedModel valide (contrairement à son prédécesseur hub.Module pour TF1).

Rubrique avancée : à quoi s'attendre de SavedModel après le chargement

Selon le contenu du SavedModel, le résultat de obj = hub.load(...) peut être appelé de différentes manières (comme expliqué de manière beaucoup plus détaillée dans le Guide SavedModel de TensorFlow :

  • Les signatures de service du SavedModel (le cas échéant) sont représentées sous la forme d'un dictionnaire de fonctions concrètes et peuvent être appelées comme tensors_out = obj.signatures["serving_default"](**tensors_in) , avec des dictionnaires de tenseurs codés par l'entrée et la sortie respectives noms et sous réserve des contraintes de forme et de dtype de la signature.

  • Les méthodes @tf.function décorées de l'objet enregistré (le cas échéant) sont restaurées en tant qu'objets tf.function qui peuvent être appelés par toutes les combinaisons d'arguments Tensor et non-Tensor pour lesquels la tf.function a été tracée avant l'enregistrement. En particulier, s'il existe une méthode obj.__call__ avec des traces appropriées, obj lui-même peut être appelé comme une fonction Python. Un exemple simple pourrait ressembler à output_tensor = obj(input_tensor, training=False) .

Cela laisse une énorme liberté dans les interfaces que SavedModels peut implémenter. L' interface Reusable SavedModels pour obj établit des conventions telles que le code client, y compris les adaptateurs tels que hub.KerasLayer , sache comment utiliser le SavedModel.

Certains SavedModels peuvent ne pas suivre cette convention, en particulier des modèles entiers non destinés à être réutilisés dans des modèles plus grands, et fournir uniquement des signatures de service.

Les variables entraînables dans un SavedModel sont rechargées comme entraînables, et tf.GradientTape les regardera par défaut. Voir la section sur le réglage fin ci-dessous pour quelques mises en garde, et envisagez d'éviter cela pour les débutants. Même si vous souhaitez affiner, vous voudrez peut-être voir si obj.trainable_variables conseille de ne réentraîner qu'un sous-ensemble des variables initialement entraînables.

Création de modèles enregistrés pour TF Hub

Aperçu

SavedModel est le format de sérialisation standard de TensorFlow pour les modèles entraînés ou les pièces de modèle. Il stocke les poids entraînés du modèle avec les opérations TensorFlow exactes pour effectuer son calcul. Il peut être utilisé indépendamment du code qui l'a créé. En particulier, il peut être réutilisé dans différentes API de création de modèles de haut niveau comme Keras, car les opérations TensorFlow sont leur langage de base commun.

Sauvegarde depuis Keras

À partir de TensorFlow 2, tf.keras.Model.save() et tf.keras.models.save_model() utilisent par défaut le format SavedModel (pas HDF5). Les SavedModels résultants qui peuvent être utilisés avec hub.load() , hub.KerasLayer et des adaptateurs similaires pour d'autres API de haut niveau dès qu'ils sont disponibles.

Pour partager un modèle Keras complet, enregistrez-le simplement avec include_optimizer=False .

Pour partager une partie d'un modèle Keras, faites de la pièce un modèle en soi, puis enregistrez-le. Vous pouvez soit présenter le code comme ça depuis le début...

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

... ou découpez la pièce à partager après coup (si elle s'aligne sur la superposition de votre modèle complet) :

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 sur GitHub utilise la première approche pour BERT (voir nlp/tools/export_tfhub_lib.py , notez la séparation entre core_model pour l'exportation et le pretrainer pour restaurer le point de contrôle) et la dernière approche pour ResNet (voir legacy/image_classification/tfhub_export.py ).

Enregistrement à partir de TensorFlow de bas niveau

Cela nécessite une bonne connaissance du guide SavedModel de TensorFlow.

Si vous souhaitez fournir plus qu'une simple signature de diffusion, vous devez implémenter l' interface Reusable SavedModel . Conceptuellement, cela ressemble à

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

Réglage fin

La formation des variables déjà formées d'un SavedModel importé avec celles du modèle qui l'entoure s'appelle affiner le SavedModel. Cela peut entraîner une meilleure qualité, mais rend souvent la formation plus exigeante (peut prendre plus de temps, dépendre davantage de l'optimiseur et de ses hyperparamètres, augmenter le risque de surajustement et nécessiter une augmentation des ensembles de données, en particulier pour les CNN). Nous conseillons aux consommateurs de SavedModel de ne se pencher sur les ajustements qu'après avoir établi un bon régime d'entraînement, et uniquement si l'éditeur de SavedModel le recommande.

Le réglage fin modifie les paramètres de modèle "continus" qui sont formés. Il ne modifie pas les transformations codées en dur, telles que la tokenisation de l'entrée de texte et le mappage des jetons à leurs entrées correspondantes dans une matrice d'intégration.

Pour les consommateurs SavedModel

Création d'un hub.KerasLayer comme

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

permet d'affiner le SavedModel chargé par la couche. Il ajoute les poids entraînables et les régularisateurs de poids déclarés dans le SavedModel au modèle Keras, et exécute le calcul du SavedModel en mode apprentissage (pensez à l'abandon, etc.).

Le colab de classification d'images contient un exemple de bout en bout avec un réglage fin facultatif.

Réexporter le résultat du réglage fin

Les utilisateurs avancés peuvent souhaiter enregistrer les résultats du réglage fin dans un modèle enregistré qui peut être utilisé à la place de celui chargé à l'origine. Cela peut être fait avec un code comme

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)

Pour les créateurs de SavedModel

Lors de la création d'un SavedModel pour le partage sur TensorFlow Hub, pensez à l'avance si et comment ses utilisateurs doivent l'ajuster, et fournissez des conseils dans la documentation.

La sauvegarde à partir d'un modèle Keras devrait permettre à tous les mécanismes de réglage fin de fonctionner (sauvegarder les pertes de régularisation du poids, déclarer des variables entraînables, tracer __call__ pour training=True et training=False , etc.)

Choisissez une interface de modèle qui fonctionne bien avec le flux de gradient, par exemple, les logits de sortie au lieu des probabilités softmax ou des prédictions top-k.

Si le modèle utilise l'abandon, la normalisation par lots ou des techniques de formation similaires qui impliquent des hyperparamètres, définissez-les sur des valeurs qui ont du sens pour de nombreux problèmes cibles et tailles de lots attendus. (Au moment d'écrire ces lignes, économiser de Keras ne permet pas aux consommateurs de les ajuster facilement.)

Les régularisateurs de poids sur les couches individuelles sont enregistrés (avec leurs coefficients de force de régularisation), mais la régularisation du poids à partir de l'optimiseur (comme tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) est perdue. Informez les consommateurs de votre SavedModel en conséquence.