PyTorch-Modelle auf Cloud TPU trainieren

In diesem Dokument werden die Details zum Ausführen von PyTorch-Modellen auf Cloud TPUs beschrieben. Wie Sie PyTorch-Modelle auf Cloud TPU trainieren, hängt davon ab, welche Cloud TPU-Architektur Sie verwenden. Weitere Informationen finden Sie unter Systemarchitektur.

Cloud TPU-VMs

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

PyTorch/XLA erstellt einen neuen lokalen TensorFlow-Server, wenn ein neues PyTorch/XLA-Programm ausgeführt wird. Der XRT (XLA Runtime-Client) stellt eine Verbindung zum lokalen TensorFlow-Server her.

Bild

Weitere Informationen zum Trainieren eines PyTorch-Modells auf TPU-VMs finden Sie unter Kurzanleitung: Cloud TPU VM PyTorch.

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. Auf dem TPU-Knoten wird ein TensorFlow-Server ausgeführt. Der XRT (XLA Runtime-Client) stellt über das Netzwerk eine Verbindung zum TensorFlow-Server her.

Bild