PyTorch integration

This page explains Vertex AI's PyTorch integration and provides resources that show you how to use PyTorch on Vertex AI. Vertex AI's PyTorch integration makes it easier for you to train, deploy, and orchestrate PyTorch models in production.

Run code in notebooks

Vertex AI provides two options for running your code in notebooks, Colab Enterprise and Vertex AI Workbench. To learn more about these options, see choose a notebook solution.

Prebuilt containers for training

Vertex AI provides prebuilt Docker container images for model training. These containers are organized by machine learning frameworks and framework versions and include common dependencies that you might want to use in your training code. To learn about which PyTorch versions have prebuilt training containers and how to train models with a prebuilt training container, see Prebuilt containers for custom training.

Prebuilt containers for serving predictions

Vertex AI provides prebuilt Docker container images for serving both batch and online predictions. These containers are organized by machine learning frameworks and framework versions and include common dependencies that you might want to use in your prediction code. To learn about which PyTorch versions have prebuilt prediction containers and how to serve models with a prebuilt prediction container, see Prebuilt containers for custom training.

Distributed training

You can run distributed training of PyTorch models on Vertex AI. For multi-worker training, you can use Reduction Server to optimize performance even further for all-reduce collective operations. To learn more about distributed training on Vertex AI, see Distributed training.

Resources for using PyTorch on Vertex AI

To learn more and start using PyTorch in Vertex AI, see the following resources:

What's next

  • Tutorial: Use Vertex AI to train a PyTorch image classification model in one of Vertex AI's prebuilt container environments by using the Google Cloud console.

    To follow step-by-step guidance for this task directly in the Google Cloud console, click Guide me:

    Guide me