Deploy models in PyTorch with Torchserve 🚀 | by Francesco Zuppichini | Towards Data Science | Towards Data Science

All the code used in this article is here
Recently, PyTorch has introduced its new production framework to properly serve models, called torchserve.So, without further due, let’s present today’s roadmap:
Installation with DockerExport your modelDefine a handlerServe our model
Installation with Docker
Export your model
Define a handler
Serve our model
To showcase torchserve, we will serve a fully trained ResNet34 to perform image classification.
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
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_classifierobject_detectortext_classifierimage_segmenter
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 subclassBaseHandler . 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!
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:
load the model
create a dummy input
trace the input through the model using torch.jit.trace
save the model
Create the .mar file
Official doc here
You need to install torch-model-archiver
git clone https://github.com/pytorch/serve.gitcd serve/model-archiverpip 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 .
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! 🥳
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
towardsdatascience.com
towardsdatascience.com
Thank you for reading.