Understanding TorchServe’s BaseHandler
This is the second post of a series, where we try to get familiar with TorchServe. In the first post, we learned the workflow of deploying a model using TorchServe given a default setup. The main step is to create a “.mar” file using the command torch-model-archiver. This file contains all the information we need for deploying our model. One file we need to include in this “.mar” file is the handler, which contains all the information about preprocessing, inference, and postprocessing that we need to make predictions. In the previous post, we used the default handler ImageClassifier, which can be used for image classification problems that are trained on ImageNet and predict at least 5 classes. For more details, have a look here:
However, in the majority of applications, we need to define a custom handler script. TorchServe offers a class called BaseHandler, where the steps of preprocessing, inference, and postprocessing are handled for a default setup. For deploying our own models, we can create our own handler class, which inherits from the BaseHandler. We only change the necessary parts. For that, this post is only dedicated to the BaseHandler class and to understanding what it does. The complete code for the BaseHandler can be found on GitHub. Here, we will only look at the parts of it, that we need to customize. This is how the class is structured:
The class contains more methods, but usually, we only need to care about the preprocess, inference, and postprocess method, and of course initialize additional variables in the __init__ method. Let’s have a look at the above metods in a bit more detail.
__init__
As always in this part, the needed variables are defined. Some important are: model, mapping, and context:
Here we might additionally add transformations or whatever we need for our model.
initialize
This method loads the model.pt file and sets it to eval-mode. First, it tries to load torchscript else it loads the state_dict-based model. As input, it gets the context-json file with all the needed information to load the model, as e.g its location. It raises an error when the model.py file is missing. It also sets the device and map_location variables.
preprocess
In this method, the preprocessing steps are defined, which need to be performed on the data, before applying the model. As input, it gets the data as a list and it returns the preprocessed data as a tensor. These steps need to be customized!
inference
In this method, the model is applied and the actual predictions are created. It receives the data as returned from the preprocess-function. This should match the model input shape. The predictions are returned as a tensor.
postprocess
This method receives as input data the predictions returned from the inference-function. It then creates the actual output of the model as a list. These steps need to be customized!
👉 In the next post, we will take this BaseHandler and customize it for the MNIST dataset to predict hand-written digits. 👀
Further Reading
Find more Data Science and Machine Learning posts here:
