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 de 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. Deuxièmement, les modèles peuvent être directement lus depuis le stockage distant vers TensorFlow. Quoi qu'il en soit, les appels aux fonctions tensorflow_hub dans le code Python réel peuvent et doivent continuer à utiliser les URL canoniques tfhub.dev des modèles, qui sont portables entre les systèmes et navigables pour la documentation. Dans les rares cas où le code utilisateur a besoin de l'emplacement réel du système de fichiers (après téléchargement et décompression, ou après résolution d'un descripteur de modèle en chemin de système de fichiers), il peut être obtenu par la fonction hub.resolve(handle) .

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

La bibliothèque tensorflow_hub met par défaut en cache les modèles sur le système de fichiers lorsqu'ils ont été téléchargés depuis tfhub.dev (ou d'autres sites d'hébergement ) et décompressés. Ce mode est recommandé pour la plupart des environnements, sauf si l'espace disque est rare mais que la bande passante et la latence du réseau sont superbes.

L'emplacement de téléchargement est par défaut un répertoire temporaire local mais peut être personnalisé en définissant la variable d'environnement TFHUB_CACHE_DIR (recommandé) ou en passant l'indicateur de ligne de commande --tfhub_cache_dir . L'emplacement de cache par défaut /tmp/tfhub_modules (ou quel que soit l'évaluation os.path.join(tempfile.gettempdir(), "tfhub_modules") ) devrait fonctionner dans la plupart des cas.

Les utilisateurs qui préfèrent la mise en cache persistante lors des redémarrages 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 celle-ci à ~/.bashrc

export TFHUB_CACHE_DIR=$HOME/.cache/tfhub_modules

... redémarrez le shell, puis cet emplacement sera utilisé. Lorsque vous utilisez un emplacement persistant, sachez qu'il n'y a pas de nettoyage automatique.

Lecture à partir du stockage distant

Les utilisateurs peuvent demander à la bibliothèque tensorflow_hub de lire directement les modèles à partir du stockage distant (GCS) au lieu de télécharger les modèles localement avec

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

ou en définissant l'indicateur de ligne de commande --tfhub_model_load_format sur UNCOMPRESSED . De cette façon, aucun répertoire de mise en cache n'est nécessaire, ce qui est particulièrement utile dans les environnements offrant peu d'espace disque mais 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 entrera en conflit avec le runtime TPU puisque la charge de travail de calcul est déléguée à une autre machine qui n'a pas accès à l'emplacement du cache par défaut. Il existe deux solutions de contournement pour cette situation :

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

La solution la plus simple consiste à demander à la bibliothèque tensorflow_hub de lire les modèles du compartiment GCS de TF Hub comme expliqué ci-dessus. Les utilisateurs disposant de leur propre compartiment GCS peuvent à la place spécifier un répertoire dans leur compartiment comme emplacement de cache avec un code tel que

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 : Pour plus d'informations sur les handles valides, cliquez ici .