Use TorchServe with a customized handler script
Introduction
In this third post in our TorchServe series, we deploy a model to predict hand-written digits using the MNIST dataset. The full example code is on GitHub. This repository also includes a simple web application of the model, which we will cover in the next post of this series. For serving the model we follow the workflow from the first post of this series:
We customize the handler class as explained in the second post of this series:
We use a modified version of an example from the TorchServe GitHub repository. The example given there deploys a model for the MNIST dataset with a customized handler script. However, the handler script inherits from the ImageClassifier and not from the BaseHandler. Here, to keep it more general we on purpose customize the BaseHandler, which can be used for all kind of models and not only for image classification problems.
Download the Model
First. letās download the model script and its serialized file, we want to deploy.
wget https://github.com/pytorch/serve/tree/master/examples/image_classifier/mnist/model.py
wget https://github.com/pytorch/serve/tree/master/examples/image_classifier/mnist/mnist_cnn.ptThe model script can be found here. It is a simple CNN, important for us is that the images have been normalized before training it. This needs to be considered in the preprocessing of the handler script.
Customized Handler Script
The main components of the handler script are the preprocess, inference and postprocess methods. We only need to customize the pre- and postprocess methods and the initial initialization. This is what the code looks like:
import io
from ts.torch_handler.base_handler import BaseHandler
from torchvision import transforms
from PIL import Image
from mnist import Net
import torch
import json
import os
class MNISTHandler(BaseHandler):
def __init__(self):
super(MNISTHandler, self).__init__()
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
def preprocess_one_image(self, req):
"""
Process one single image.
"""
# get image from the request
image = req.get("data")
if image is None:
image = req.get("body")
# create a stream from the encoded image
image = Image.open(io.BytesIO(image))
image = self.transform(image)
# add batch dim
image = image.unsqueeze(0)
return image
def preprocess(self, requests):
"""
Process all the images from the requests and batch them in a Tensor.
"""
images = [self.preprocess_one_image(req) for req in requests]
images = torch.cat(images)
return images
def postprocess(self, data):
"""The post process of MNIST converts the predicted output response to a label.
Args:
data (list): The predicted output from the Inference with probabilities is passed
to the post-process function
Returns:
list : A list of dictionaries with predictions and explanations is returned
"""
return data.argmax(1).tolist()Letās go through the customized parts:
- __init__ ā We can inherit the init method from the BaseHandler, but we add the definition of the transforms (i.e. the normalization) to it.
- Preprocess ā The preprocess method gets the data as a list. We then go through this list, read each image, transform it, and put it into the correct shape. The list of images is then concatenated to a tensor. This procedure I copied from Francesco Zuppichiniās post about TorchServe.
- Postprocess ā The model predictions are the prediction probabilities of each of the 10 digits. The final prediction is the index of the maximal prediction probability, which is the predicted number. This has to be returned as a list.
Create the archived (.mar) File
As explained in the first post of this series, we need to create a ā.marā-file, which contains all the information for deploying the model. More details can be found in the TorchServe GitHub repository. For this case the command looks like this:
torch-model-archiver --model-name mnist --version 2.0 \
--model-file mnist.py \
--serialized-file mnist_cnn.pt \
--handler mnist_handler_base.pyNow, we create a folder called āmodel-storeā and move the just created mnist.mar file into it.
mv mnist.mar model-store/Make Predictions
We again follow the same steps as in the first post of this series. To make predictions, we deploy the model using Docker. We use the latest TorchServe image (āpytorch/torchserve:latestā) and directly start TorchServe. The folder, where the ā.marā-file is stored (āmodel-storeā) is included as a volume to the container. More details and information about customized images can also be found on the TorchServe GitHub repository
docker run --rm -it -p 8080:8080 -p 8081:8081 --name mar \
-v $(pwd)/model-store:/home/model-server/model-store \
pytorch/torchserve:latest torchserve \
--start --model-store model-store
--models mnist=mnist.marOpen a new terminal and check the available models:
curl http://127.0.0.1:8081/models/Output:
{
"models": [
{
"modelName": "mnist",
"modelUrl": "mnist.mar"
}
]
}Test images can be downloaded from the TorchServe GitHub Repository, e.g.:
wget https://github.com/pytorch/serve/tree/master/examples/image_classifier/mnist/test_data/3.pngWe use this example image to make predictions:
curl http://127.0.0.1:8080/predictions/mnist -T 3.pngI hope this helps you to get started with your own project. Of course MNIST is a very simple example, but the concepts and workflow remains the same for all custom model deployments.
Further Readings
- https://github.com/pytorch/serve/blob/master/ts/torch_handler/base_handler.py
- https://pytorch.org/serve/default_handlers.html
- https://towardsdatascience.com/deploy-models-and-create-custom-handlers-in-torchserve-fc2d048fbe91
- https://github.com/pytorch/serve/blob/master/model-archiver/README.md
Find more Data Science and Machine Learning posts here:
