Mise en cache des téléchargements de modèles depuis TF Hub

Aperçu

La bibliothèque tensorflow_hub prend actuellement en charge deux modes de téléchargement des modèles. Par défaut, un modèle est téléchargé sous forme d'archive compressée et mis en cache sur le disque. Il est également possible de lire directement les modèles depuis un stockage distant dans TensorFlow. Dans les deux cas, les appels aux fonctions tensorflow_hub dans le code Python peuvent et doivent continuer à utiliser les URL canoniques tfhub.dev des modèles, qui sont portables d'un système à l'autre et permettent de naviguer dans la documentation. Dans les rares cas où le code utilisateur a besoin de l'emplacement réel du modèle sur le système de fichiers (après téléchargement et décompression, ou après résolution d'un descripteur de modèle en un chemin d'accès au système de fichiers), celui-ci peut être obtenu grâce à la fonction hub.resolve(handle) .

Mise en cache des téléchargements compressés

Par défaut, la bibliothèque tensorflow_hub met en cache les modèles sur le système de fichiers après leur téléchargement depuis tfhub.dev (ou d'autres sites d'hébergement ) et leur décompression. Ce mode est recommandé dans la plupart des environnements, sauf en cas d'espace disque limité mais de bande passante et de latence réseau optimales.

Par défaut, le répertoire de téléchargement est un répertoire temporaire local, mais il est possible de le personnaliser en définissant la variable d'environnement TFHUB_CACHE_DIR (recommandé) ou en utilisant l'option de ligne de commande --tfhub_cache_dir . L'emplacement de cache par défaut /tmp/tfhub_modules (ou le résultat de l'évaluation de os.path.join(tempfile.gettempdir(), "tfhub_modules") devrait convenir dans la plupart des cas.

Les utilisateurs qui préfèrent une mise en cache persistante après un redémarrage du système peuvent définir TFHUB_CACHE_DIR sur un emplacement de leur répertoire personnel. Par exemple, un utilisateur du shell bash sur un système Linux peut ajouter une ligne comme la suivante à ~/.bashrc :

export TFHUB_CACHE_DIR=$HOME/.cache/tfhub_modules

Redémarrez le terminal ; cet emplacement sera alors utilisé. Notez que l’utilisation d’un emplacement persistant ne permet pas un nettoyage automatique.

Lecture depuis un stockage distant

Les utilisateurs peuvent indiquer à la bibliothèque tensorflow_hub de lire directement les modèles depuis un stockage distant (GCS) au lieu de les télécharger localement :

os.environ["TFHUB_MODEL_LOAD_FORMAT"] = "UNCOMPRESSED"

ou en définissant l'option de ligne de commande --tfhub_model_load_format sur UNCOMPRESSED . De cette façon, aucun répertoire de cache n'est nécessaire, ce qui est particulièrement utile dans les environnements disposant de peu d'espace disque mais d'une connexion Internet rapide.

Exécution sur TPU dans les notebooks Colab

Sur colab.research.google.com , le téléchargement de modèles compressés entre en conflit avec l'environnement d'exécution TPU, car la charge de calcul est déléguée à une autre machine qui n'a pas accès par défaut à l'emplacement du cache. Deux solutions de contournement existent :

1) Utilisez un bucket GCS auquel le worker TPU peut accéder.

La solution la plus simple consiste à indiquer à la bibliothèque tensorflow_hub de lire les modèles depuis le bucket GCS de TF Hub, comme expliqué précédemment. Les utilisateurs disposant de leur propre bucket GCS peuvent également spécifier un répertoire de leur bucket comme emplacement du cache, avec un code similaire à celui-ci :

import os
os.environ["TFHUB_CACHE_DIR"] = "gs://my-bucket/tfhub-modules-cache"

...avant d'appeler la bibliothèque tensorflow_hub .

2) Rediriger toutes les lectures via l'hôte Colab

Une autre solution consiste à rediriger toutes les lectures (même celles des variables volumineuses) via l'hôte Colab :

load_options =
tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')
reloaded_model = hub.load("https://tfhub.dev/...", options=load_options)

Remarque : Consultez ici plus d'informations concernant les identifiants valides.