PyTorch: Quick Start

In this tutorial, we are going to deploy a PyTorch image classification model to Model Zoo and use it to make a prediction. The Model Zoo client library relies on the ONNX format to serialize models.

You can follow along this tutorial in any Python environment you’re comfortable with, such as a Python IDE, Jupyter notebook, or a Python terminal. The easiest option is to open this tutorial directly in colab:

Open In Colab


Install the Model Zoo client library via pip:

!pip install modelzoo-client[torch]

To deploy and use your own models, you’ll need to create an account and configure an API key. You can do so from the command line:

!modelzoo auth


First, we’ll need a PyTorch model (nn.Module) to deploy. For the sake of this quickstart, we’ll use the torchvision package to load a pretrained Mobile Net V2 image classification model.

import torch
import torchvision

model = torchvision.models.mobilenet_v2(pretrained=True)


To deploy this pipeline to a production-ready HTTP endpoint, use the modelzoo.torch.deploy() function. Since the Model Zoo PyTorch client library relies on the ONNX format for serialization. modelzoo.torch.deploy() uses the same arguments as torch.onnx.export(), except that it does not require a filename. Instead, the ONNX-serialized model will be directly uploaded to Model Zoo.

import modelzoo.torch

model_name = modelzoo.torch.deploy(
    model, torch.randn(1, 3, 224, 224), input_names=["input"], output_names=["output"]

That’s all there is to it! Behind the scenes, Model Zoo serialized your model to ONNX, uploaded it to object storage, deployed a container to serve any HTTP requests made to the model, and set up a load balancer to route requests to multiple model shards. If you’d like, take some time to explore the model via the Web UI link. There you’ll be able to modify documentation, test the model with raw or visual inputs, monitor metrics and/or logs. By default, only your account (or anybody you share your API key with) will be able to access this model.

You can specify the name of the model you’d like to deploy via a model_name argument. If a name is omitted, Model Zoo will choose a unique one for you. Model names must be unique to your account.


First, let’s load a picture of a dog to use as a test for the image classification model. We’ll also load the JSON metadata that maps the class indices to human-readable labels.

from PIL import Image
from io import BytesIO
import requests

image_response = requests.get(
input_image =

class_to_labels = requests.get(

Next, we’ll use our Python client library to query the model for a prediction. modelzoo.torch.predict() requires the model_name and a Python dictionary mapping input layer string names to a dictionary in onnx.TensorProto format. In this example, we’ll also transform the input image and apply the appropriate transformations as specified by the model documentation.

import modelzoo.torch
import numpy as np

import onnx
from torchvision import transforms

preprocess = transforms.Compose(
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
input_tensor = preprocess(input_image)
input_tensor = input_tensor.unsqueeze(0)

full_output = modelzoo.torch.predict(
        "input": {
            "dims": input_tensor.shape,
            "data_type": onnx.TensorProto.DataType.FLOAT,
            "float_data": input_tensor.numpy().tolist(),

output = full_output["outputs"]["output"]["floatData"]
model_prediction_idx = str(np.argmax(output))
print("Model prediction: {}".format(class_to_labels[model_prediction_idx]))

Great! At this point, we’ve used our image classification model to successfully make a prediction on a new image.


By default, Model Zoo will deploy your model and wait for it to get into a HEALTHY state, meaning that it’s ready for predictions. You can always check on the state of a model by using the function:

To save resources and shut down any model if you aren’t using it, you can use modelzoo.stop():


With Model Zoo you can manage model state manually, or automatically. By default, our free trial will stop any model where there has been no request activity for 15 minutes, saving you resources if you forget to stop manually. Our unlimited version has more options for controlling autoscaling behavior.

Interested in what you’ve seen and want to test drive an unlimited version of Model Zoo? Apply to our private beta and reach out at to learn more.