Utilisation de TPU avec PyTorch
J'essaie d'utiliser le TPU de Google Cloud de Colab. J'ai pu le faire en suivant le tutoriel en utilisant Tensorflow.
Quelqu'un sait-il s'il est possible d'utiliser les TPU à l'aide de PyTorch? Si oui, comment puis-je le faire? Avez-vous un exemple?
Consultez notre référentiel pytorch/xla où vous pouvez commencer à former des modèles PyTorch sur des TPU.
En outre, vous pouvez même utiliser des TPU gratuits sur Colab avec PyTorch avec ces blocs-notes Colab .
Actuellement, il n'est pas possible d'utiliser Cloud TPU avec PyTorch car il est conçu spécifiquement pour Tensorflow.
Mais, selon cette nouvelle sur les produits publiée il y a trois jours sur le blog de Google Cloud , "les ingénieurs de l'équipe TPU de Google collaborent activement avec les principaux développeurs PyTorch pour connecter PyTorch aux Cloud TPU".
À ce jour, PyTorch Lightning permet d'exécuter trivialement le code PyTorch sur les TPU (vous aurez besoin de la bibliothèque XLA installée). Depuis leur carnet de démo sur colab :
from pytorch_lightning import Trainer
model = CoolSystem()
# most basic trainer, uses good defaults
trainer = Trainer(num_tpu_cores=8)
trainer.fit(model)