Cambio di versione software su Cloud TPU

Panoramica

La versione software del framework in esecuzione sulla TPU deve corrispondere alla versione in esecuzione sulla VM locale. Questa versione del software può essere cambiata in una Cloud TPU in esecuzione, senza eliminare e ricreare la TPU. Ciò consente anche di configurare Cloud TPU con versioni notturne specifiche di framework software. È comunque consigliabile selezionare una versione supportata di questi framework.

Utilizzo

Il metodo consigliato per cambiare versione è utilizzare la libreria Python cloud-tpu-client.

Esempio di utilizzo per TensorFlow.

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

Questa configurazione configura Cloud TPU in modo che corrisponda alla versione di TensorFlow in esecuzione sulla tua VM locale, incluse le release ufficiali e le build notturne con data.

Il parametro restart_type dell'API configure_tpu_version definisce il comportamento di riavvio della TPU quando cambi versione. Le opzioni sono 'sempre' (impostazione predefinita) e 'se necessario'.

  • 'sempre' può essere utilizzato per correggere una TPU con, ad esempio, lo stato UNHEALTHY_TENSORFLOW o che restituisce errori di memoria fughe (OOM) a causa di risorse che sono state divulgate da un'esecuzione precedente. Se questa opzione è impostata, la TPU viene riavviata anche quando non è installata una nuova versione del framework.

  • 'ifneeded' può essere utile perché non riavvia il runtime se è già nella versione corretta, quindi non aggiungerà alcun tempo di avvio significativo a uno script di addestramento. Se questa opzione è impostata, la TPU viene riavviata solo se non è installata la versione del framework corretta.

La libreria comunica direttamente con Cloud TPU, quindi questo codice deve essere eseguito in una VM nella stessa rete. Ti consigliamo di eseguirlo all'interno del codice per il resto del modello.

Opzioni software aggiuntive

TensorFlow include una stringa tf.__version__, il modo più semplice per configurare la versione corretta. Altre opzioni software includono:

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

Ad esempio, per configurare una TPU da eseguire con l'ultima build notturna di PyTorch.

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