Deploy models and create custom handlers in Torchserve

Deploy models and create custom handlers in Torchserve :rocket:

let’s put a model into production

Francesco Zuppichini

Follow

Jun 11 ·6min read

All the code used in this article is here

If you are here it means you are interested in torchserve , the new tool to properly put models into production. So, without further due, let’s present today’s roadmap:

1. Installation with Docker
2. Export your model
3. Define a handler
4. Serve our model

To showcase torchserver, we will serve a fully trained ResNet34 to perform image classification.

Installation with Docker

Official doc here

The best way to install torchserve is with docker. You just need to pull the image.

You can use the following command to save the latest image.

docker pull pytorch/torchserve:latest

All the tags are available here

More about docker and torchserve here

Handlers

Official doc here

Handlers are the ones responsible to make a prediction using your model from one or more HTTP requests.

Default handlers

Torchserve supports the following default handlers

image_classifier
object_detector
text_classifier
image_segmenter

But keep in mind that none of them supports batching requests!

Custom handlers

torchserve exposes a rich interface to do almost everything you want. An Handler is just a class that must have three functions

• preprocess
• inference
• postprocess

You can create your own class or just subclass BaseHandler . The main advantage of subclasssing BaseHandler is to have the model loaded accessible at self.model . The following snippet shows how to subclass BaseHandler

Going back to our image classification example. We need to

• get the images from each request and preprocess them
• get the prediction from the model
• send back a response

Preprocess

The .preprocess function takes an array of requests. Assuming we are sending an image to the server, the serialized image can be accessed from the data or body field of the request. Thus, we can just iterate over all requests and preprocess individually each image. The full code is shown below.

self.transform is our preprocess transformation, nothing fancy. This is a classic preprocessing step for models trained on ImageNet.

After we have preprocessed each image in each request we concatenate them to create a pytorch Tensor.

Inference

This step is very easy, we get the tensor from the .preprocess function and we extract the prediction for each image.

Postprocess

Now we have our predictions for each image, we need to return something to the client. Torchserve always expects an array to be returned. BaseHandler also automatically opens a .json file with the mapping index -> label (we are going to see it later how to provide such file) and store it at self.mapping . We can return an array of dictionaries with the label and index class for each prediction

Wrapping everything together, our glorious handler looks like

Since all the handling logic encapsulated in a class, you can easily unit test it!

Export your model

Official doc here

Torchserve expects a .mar file to be provided. In a nutshell, the file is just your model and all the dependencies packed together. To create one need to first export our trained model.

Export the model

There are three ways to export your model for torchserve. The best way that I have found so far is to trace the model and store the results. By doing so we do not need to add any additional files to torchserve.

Let’s see an example, we are going to deploy a fully trained ResNet34 model.

In order, we:

torch.jit.trace

Create the .mar file

Official doc here

You need to install torch-model-archiver

git clone https://github.com/pytorch/serve.git
cd serve/model-archiver
pip install .

Then, we are ready to create the .mar file by using the following command

torch-model-archiver --model-name resnet34 \--version 1.0 \--serialized-file resnet34.pt \--extra-files ./index_to_name.json,./MyHandler.py \--handler my_handler.py  \--export-path model-store -f

In order. The variable --model-name defines the final name of our model. This is very important since it will be the namespace of the endpoint that will be responsible for its predictions. You can also specify a --version . --serialized-file points to the stored .pt model we created before. --handler is a python file where we call our custom handler. In general, it always looks like this:

It exposes a handle function from which we call the methods in the custom handler. You can use the default names to use the default handled (e.g. --handler image_classifier ).

In --extra-files you need to pass the path to all the files your handlers are using. In our case, we have to add the path to the .json file with all the human-readable labels names and MyHandler.py file in which we have the class definition for MyHandler.

One minor thing, if you pass an index_to_name.json file, it will be automatically loaded into the handler and be accessible at self.mapping .

--export-path is where the .mar file will be stored, I also added the -f to overwrite everything in it.

If everything went smooth, you should see resnet34.mar stored into ./model-store .

Serve our model

This is an easy step, we can run the torchserve docker container with all the required parameters

docker run --rm -it \-p 3000:8080 -p 3001:8081 \-v $(pwd)/model-store:/home/model-server/model-store pytorch/torchserve:0.1-cpu \torchserve --start --model-store model-store --models resnet34=resnet34.mar

I am binding the container port 8080 and 8081 to 3000 and 3001 respectively (8080/8081 were already in used in my machine). Then, I am creating a volume from ./model-store (where we stored the .mar file) to the container default model-store folder. Lastly, I am calling torchserve by padding the model-store path and a list of key-value pairs in which we specify the model name for each .mar file.

At this point, torchserve has one endpoint /predictions/resnet34 to which we can get a prediction by sending an image. This can be done using curl

curl -X POST http://127.0.0.1:3000/predictions/resnet34 -T inputs/kitten.jpg

The response

{
  "label": "tiger_cat",
  "index": 282
}

It worked!

Summary

To recap, in this article we have covered:

• torchserve installation with docker
• default and custom handlers
• model archive generation
• serving the final model with docker

All the code is here

If you like this article and pytorch, you may also be interested in these my other articles

Thank you for reading.

Francesco