Making a prediction with Sagemaker PyTorch

As a data scientist or software engineer, one of the most important tasks that you might have to perform is making accurate predictions from your data. PyTorch is a popular Python-based deep learning framework that is widely used for creating and training neural networks. Amazon SageMaker is a fully managed machine learning service that makes it easy to build, train, and deploy machine learning models at scale. In this blog post, we will explore how to use SageMaker PyTorch to make predictions on a dataset.

Table of Contents

1. What is PyTorch?
2. What is SageMaker?
3. How to make a prediction with SageMaker PyTorch
   • Step 1: Prepare your dataset
   • Step 2: Train your model
   • Step 3: Deploy your model
   • Step 4: Make a prediction
4. Conclusion

What is PyTorch?

PyTorch is an open-source machine learning library that is used for building and training neural networks. It is developed primarily by Facebook’s AI research group and is based on the Torch library. PyTorch has become very popular among data scientists and machine learning engineers because it is easy to use and highly flexible.

What is SageMaker?

Amazon SageMaker is a fully managed machine learning service that enables developers and data scientists to build, train, and deploy machine learning models at scale. SageMaker provides a fully integrated development environment (IDE) for building, training, and deploying machine learning models. It also provides pre-built algorithms and frameworks, such as PyTorch, to make it easy to get started with machine learning.

How to make a prediction with SageMaker PyTorch

To make a prediction with SageMaker PyTorch, you will need to follow these steps:

Step 1: Prepare your dataset

Before you can make a prediction, you need to prepare your dataset. This involves cleaning and preprocessing your data, splitting it into training and testing sets, and converting it into a format that can be used by PyTorch. In Sagemaker, we have to specify to the PyTorch estimator the location of the training and testing data, whether it’s a path to an S3 bucket or a local file system path if utilizing local mode. In this example, we retrieve the MNIST data from a public S3 bucket and transfer it to your default bucket.

Sample code:

import logging
import boto3
from botocore.exceptions import ClientError

# Download training and testing data from a public S3 bucket


def download_from_s3(data_dir="./data", train=True):
    """
    Download MNIST dataset and convert it to numpy array
    """

    if not os.path.exists(data_dir):
        os.makedirs(data_dir)

    if train:
        images_file = "train-images-idx3-ubyte.gz"
        labels_file = "train-labels-idx1-ubyte.gz"
    else:
        images_file = "t10k-images-idx3-ubyte.gz"
        labels_file = "t10k-labels-idx1-ubyte.gz"

    # download objects
    s3 = boto3.client("s3")
    bucket = f"sagemaker-example-files-prod-{region}"
    for obj in [images_file, labels_file]:
        key = os.path.join("datasets/image/MNIST", obj)
        dest = os.path.join(data_dir, obj)
        if not os.path.exists(dest):
            s3.download_file(bucket, key, dest)
    return


download_from_s3("./data", True)
download_from_s3("./data", False)

# Upload to the default bucket

prefix = "DEMO-mnist-pytorch"
bucket = sess.default_bucket() # you can change it with your own bucket
loc = sess.upload_data(path="./data", bucket=bucket, key_prefix=prefix)

channels = {"training": loc, "testing": loc}

Step 2: Train your model

Once you have prepared your dataset, you can train your PyTorch model using SageMaker. SageMaker provides pre-built PyTorch containers that you can use to train your model. You can also customize these containers to meet your specific needs.

To train a PyTorch model, we first need to implement an entrypoint. The training entry point is a Python script encompassing all the necessary code for training a PyTorch model. This script serves as the entry point for executing the training job through the SageMaker PyTorch Estimator.

# Set local_mode to True to run the training script on the machine that runs this notebook

local_mode = False

if local_mode:
    instance_type = "local"
else:
    instance_type = "ml.c4.xlarge"

est = PyTorch(
    entry_point="train.py", # you can define this script based on your preference
    source_dir="code",  # directory of your training script
    role=role,
    framework_version="1.5.0",
    py_version="py3",
    instance_type=instance_type,
    instance_count=1,
    volume_size=250,
    output_path=output_path,
    hyperparameters={"batch-size": 128, "epochs": 1, "learning-rate": 1e-3, "log-interval": 100},
)

Step 3: Deploy your model

After you have trained your PyTorch model, you can deploy it using SageMaker. SageMaker provides several deployment options, including deploying your model as an endpoint or as a batch transform job.

To deploy a model, we need a PyTorchModel class which allows you to create an environment for model inference utilizing your model artifact. Similar to the PyTorch class elaborated in this notebook for training a PyTorch model, it serves as a high-level API for configuring a Docker image for your model hosting service.

model = PyTorchModel(
    entry_point="inference.py", # this file can be defined so it take a tensor as an input and return the prediction
    source_dir="code",
    role=role,
    model_data=pt_mnist_model_data,
    framework_version="1.5.0",
    py_version="py3",
)

To prepare for inference, create an entry point file named inference.py in the code directory. This file will include the following functions:

1. model_fn: Responsible for loading the model checkpoint from the model directory in the inference image and returning a PyTorch model. Example implementation:

def model_fn(model_dir):
    model = Net()
    with open(os.path.join(model_dir, "model.pth"), "rb") as f:
        model.load_state_dict(torch.load(f))
    model.to(device).eval()
    return model

2. input_fn: Handles data decoding for incoming requests. Signature:

def input_fn(request_body, request_content_type):

Example implementation:

def input_fn(request_body, request_content_type):
    assert request_content_type=='application/json'
    data = json.loads(request_body)['inputs']
    data = torch.tensor(data, dtype=torch.float32, device=device)
    return data

3. predict_fn: Invoked after deserialization by input_fn to make inference using the model. Signature:

def predict_fn(input_object, model):

Example implementation:

def predict_fn(input_object, model):
    with torch.no_grad():
        prediction = model(input_object)
    return prediction

4. output_fn: Performs post-processing on inference results. Signature:

def output_fn(prediction, content_type):

Example implementation:

def output_fn(predictions, content_type):
    assert content_type == 'application/json'
    res = predictions.cpu().numpy().tolist()
    return json.dumps(res)

Ensure the SageMaker PyTorch model server is invoked with data readily consumable by the model, i.e., normalized with batch and channel dimensions.

Once it is properly configured, it can be used to create a SageMaker endpoint on an EC2 instance. The SageMaker endpoint is a containerized environment that uses your trained model to make inference on incoming data via RESTful API calls.

Step 4: Make a prediction

Once your model is deployed, you can use it to make predictions on your dataset. You can do this by sending your data to the endpoint or by running a batch transform job.

Now, we proceed to test the endpoint with real MNIST test data, which we retrieve and normalize using helper functions from code.utils.

from utils.mnist import mnist_to_numpy, normalize
import random
import matplotlib.pyplot as plt

%matplotlib inline

data_dir = "/tmp/data"
X, Y = mnist_to_numpy(data_dir, train=False)

# randomly sample 16 images to inspect
mask = random.sample(range(X.shape[0]), 16)
samples = X[mask]
labels = Y[mask]

# pre-processing: convert data type, normalize, and add channel dimension
samples = normalize(samples.astype(np.float32), axis=(1, 2))
samples = np.expand_dims(samples, axis=1)

res = predictor.predict({"inputs": samples.tolist()})

The response consists of probability vectors for each sample, which we use to determine the predicted digits.

predictions = np.argmax(np.array(res, dtype=np.float32), axis=1).tolist()
print("Predicted digits: ", predictions)

Notes: All code snippets in this article are employed from AWS official github.

Conclusion

In conclusion, SageMaker PyTorch provides a powerful and flexible platform for making predictions on your data. By following the steps outlined above, you can easily prepare your dataset, train your model, deploy it, and make predictions. PyTorch’s ease of use and flexibility make it a popular choice for data scientists and machine learning engineers, and SageMaker’s fully managed service makes it easy to scale your models and deploy them in production environments.

SHARE:

About Saturn Cloud

Saturn Cloud is your all-in-one solution for data science & ML development, deployment, and data pipelines in the cloud. Spin up a notebook with 4TB of RAM, add a GPU, connect to a distributed cluster of workers, and more. Join today and get 150 hours of free compute per month.