Changement de version du logiciel sur votre Cloud TPU

Présentation

La version logicielle du framework exécutée sur votre TPU doit correspondre à la version exécutée sur votre VM locale. Cette version du logiciel peut désormais être remplacée par un Cloud TPU en cours d'exécution, sans supprimer ni recréer le TPU. Cela permet également de configurer le Cloud TPU avec des versions nocturnes spécifiques de frameworks logiciels. Il est toujours recommandé de sélectionner une version compatible de ces frameworks.

Utilisation

La méthode recommandée pour changer de version consiste à utiliser la bibliothèque Python cloud-tpu-client.

Exemple d'utilisation de TensorFlow.

import tensorflow as tf
from cloud_tpu_client import Client

c = Client()
c.configure_tpu_version(tf.__version__, restart_type='ifNeeded')

Le Cloud TPU est configuré de manière à correspondre à la version de TensorFlow exécutée sur votre VM locale, y compris les versions officielles et les versions nocturnes datées.

Le paramètre restart_type de l'API configure_tpu_version définit le comportement du redémarrage TPU lors du changement de version. Les options sont "always" (par défaut) et "ifneeded".

  • Vous pouvez utiliser "always" pour résoudre un TPU avec, par exemple, l'état UNHEALTHY_TENSORFLOW, ou qui renvoie des erreurs de mémoire saturée (OOM, Out Of Memory) en raison de fuites de ressources lors d'une exécution précédente. Si cette option est définie, le TPU est redémarré même lorsqu'une nouvelle version du framework n'est pas installée.

  • "ifneeded" peut être utile, car il ne redémarre pas l'environnement d'exécution s'il se trouve déjà à la bonne version. Il n'ajoute donc aucun temps de démarrage significatif à un script d'entraînement. Lorsque cette option est définie, le TPU ne redémarre que si la version de framework appropriée n'est pas installée.

La bibliothèque communique directement avec Cloud TPU. Ce code doit donc être exécuté sur une VM du même réseau. Il est recommandé de l'exécuter dans le code du reste de votre modèle.

Options logicielles supplémentaires

TensorFlow inclut une chaîne tf.__version__, le moyen le plus simple de configurer la version appropriée. Les autres options logicielles sont les suivantes:

  • PyTorch - pytorch-0.1, pytorch-0.5-dev20200320, pytorch-nightly
  • Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

Par exemple, pour configurer un TPU à exécuter avec la dernière version nocturne de PyTorch.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()