PyTorch-Modelle auf Cloud TPU trainieren

In diesem Dokument wird die Ausführung von PyTorch-Modellen auf Cloud TPUs beschrieben. Wie Sie PyTorch-Modelle auf Cloud TPU trainieren, hängt von der verwendeten Cloud TPU-Architektur ab. Weitere Informationen finden Sie unter Systemarchitektur.

Cloud TPU-VMs

Bei Cloud TPU-VMs gibt es keine Nutzer-VMs und Ihr Modell und Ihr Code werden auf einer VM auf dem TPU-Hostcomputer ausgeführt. Sie können eine SSH-Verbindung direkt zum TPU-Host herstellen. Sie können beliebige Codeinstallationspakete ausführen, Logs anzeigen, Code debuggen usw.

PyTorch/XLA erstellt bei jeder Ausführung eines neuen PyTorch/XLA-Programms einen lokalen TensorFlow-Server. Der XRT-Client (XLA Runtime) stellt eine Verbindung zum lokalen TensorFlow-Server her.

Image

Weitere Informationen zum Trainieren eines PyTorch-Modells auf TPU-VMs finden Sie unter Kurzanleitung für PyTorch auf Cloud TPU-VMs.

Legacy-TPU-Knoten

Wenn Sie Legacy-TPU-Knoten verwenden, benötigen Sie für jede TPU eine Nutzer-VM. Nutzer-VMs kommunizieren über gRPC-Aufrufe mit den TPU-Knoten. Der TPU-Knoten führt einen TensorFlow-Server aus. Der XRT-Client (XLA Runtime) stellt über das Netzwerk eine Verbindung zum TensorFlow-Server her.

Image