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

Le format SavedModel de TensorFlow 2 est le moyen recommandé pour partager des modèles pré-entraînés et des éléments de modèle 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 hub.load() de bas niveau et son wrapper hub.KerasLayer . (Généralement, 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 modèles existants 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 des modèles enregistrés depuis TF Hub

Utilisation d'un SavedModel dans Keras

Keras est l'API de haut niveau de TensorFlow permettant de créer des modèles d'apprentissage profond 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é-entraîné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 laboratoire de classification de texte est un exemple complet de la façon de former et d'évaluer un tel classificateur.

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

Utilisation d'un SavedModel 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 sur 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, voir Mise en cache .

Utiliser un SavedModel dans TensorFlow de bas niveau

Modèle de 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 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 de son 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).

Sujet avancé : à quoi s'attendre du SavedModel après le chargement

En fonction du contenu du SavedModel, le résultat de obj = hub.load(...) peut être invoqué de différentes manières (comme expliqué plus en détail dans le Guide SavedModel de TensorFlow :

  • Les signatures de service du SavedModel (le cas échéant) sont représentées comme 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 saisis par l'entrée et la sortie respectives. noms et soumis aux contraintes de forme et de type de la signature.

  • Les méthodes décorées @tf.function 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 réutilisable SavedModels pour obj établit des conventions telles que le code client, y compris les adaptateurs comme hub.KerasLayer , sache comment utiliser le SavedModel.

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

Les variables pouvant être entraînées dans un SavedModel sont rechargées comme pouvant être entraînées, et tf.GradientTape les surveillera par défaut. Consultez la section sur le réglage fin ci-dessous pour quelques mises en garde, et envisagez d'éviter cela pour commencer. Même si vous souhaitez affiner, vous voudrez peut-être voir si obj.trainable_variables conseille de recycler uniquement 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 ou les éléments de modèle entraînés. Il stocke les poids entraînés du modèle ainsi que 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.

Économiser de Keras

À partir de TensorFlow 2, tf.keras.Model.save() et tf.keras.models.save_model() sont par défaut au format SavedModel (et non 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-la. Vous pouvez soit présenter le code comme ça dès 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 correspond à 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)

Les modèles TensorFlow sur GitHub utilisent 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 SavedModel Guide de TensorFlow.

Si vous souhaitez fournir plus qu'une simple signature de diffusion, vous devez implémenter l' interface Réutilisable 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

L'entraînement des variables déjà entraînées d'un SavedModel importé ainsi que celles du modèle qui l'entoure est appelé réglage fin du 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 de l'ensemble de données, en particulier pour les CNN). Nous conseillons aux consommateurs de SavedModel de se pencher sur le réglage fin seulement après avoir établi un bon programme de formation, et seulement si l'éditeur SavedModel le recommande.

Le réglage fin modifie les paramètres du modèle « continus » qui sont entraînés. Il ne modifie pas les transformations codées en dur, telles que la tokenisation de la saisie de texte et le mappage des jetons avec 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 entraînement (pensez à l'abandon, etc.).

Le laboratoire 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 voudront peut-être enregistrer les résultats du réglage fin dans un SavedModel 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 à partager sur TensorFlow Hub, réfléchissez à l'avance si et comment ses consommateurs doivent le peaufiner, et fournissez des conseils dans la documentation.

L'enregistrement à partir d'un modèle Keras devrait faire fonctionner tous les mécanismes de réglage fin (économie des pertes de régularisation de poids, déclaration des variables pouvant être entraînées, traçage de __call__ pour training=True et training=False , etc.)

Choisissez une interface de modèle qui fonctionne bien avec le flux de gradient, par exemple, des 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 impliquant 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, les économies réalisées auprès de Keras ne permettent 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 de poids depuis l'optimiseur (comme tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) est perdue. Informez les consommateurs de votre SavedModel en conséquence.