How to Serve your PyTorch Models. The Newest Version of TorchServe is one... | by Dimitris Poulopoulos | Towards Data Science

You have collected your data, processed them, trained your model, fine-tuned it, and the results are promising. Where to go next? How do you make it available to the general public?
Well, if you are a TensorFlow user, you might want to fiddle with TFX. TFX is an excellent collection of tools that helps data scientists create and manage a machine learning production pipeline. But what if you are a PyTorch user? Well, you are in luck; as of April 2020, you got TorchServe.
Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!
TorchServe is a flexible and easy to use tool for serving PyTorch models. It does not have the complexity of TFX and, as a result, it does not provide that many functionalities. However, it is a straightforward way to get the job done!
TorchServe provides a set of necessary features, such as a server, a model archiver tool, an API endpoint specification, logging, metrics, batch inference and model snapshots among others. It also offers a list of advanced features, for instance, support for custom inference services, unit test and an easy way to collect benchmark data through JMeter. Currently, it is in an experimental stage, but it works like a charm for most cases. As a matter of fact, we take it for a test drive in later sections of this article.
TorchServe released its latest version (0.1.1) two days ago (June 10th 2020). It is a step forward towards GA (General Availability) and includes a set of promising new features. Among others, the newest version of TorchServe includes:
Support for HuggingFace, a deep learning library with a mission to advance and democratize NLP for everyone, with documentation and examples
Support for Nvidia Waveglow, a flow-based generative network for speech synthesis, with documentation and examples
Tight integration with Model Zoo, a deep learning registry with model archives created from popular pre-trained models
Support for AWS Cloud Formation that provides the ability to spin up a server on an EC2 instance with the convenience of a simple configuration file (YAML or JSON)
Support for Profiling TorchServe Python execution via snakevize profiler for detailed execution time reporting
Refactored documentation with clear instructions
For a detailed list of available features, read the release notes.
Installing TorchServe is straightforward. You can use the instructions provided in the documentation, but let us install it on Ubuntu using conda:
Create a new environment for your TorchServe installation (optional but recommended).
Create a new environment for your TorchServe installation (optional but recommended).
conda create -n torch python=3.8
2. Activate the new environment.
conda activate torch
3. Run the following script to install TorchServe and its dependencies. This script installs the CPU-only PyTorch, torchvision and torchtext modules in your machine. Feel free to skip any of these if you do not need them.
conda install pytorch torchvision torchtext torchserve torch-model-archiver psutil future cpuonly -c pytorch -c powerai
You are now ready to start using TorchServe and make your models available via a REST endpoint.
In this example, we deploy an MNIST classifier. We go through the steps needed, with code examples taken from TorchServe documentation.
First, we need to define a PyTorch model to solve the MNIST challenge. The code below is just an approach; copy and paste it in a file called model.py:
Then, we need to create a script that trains the model. Copy the following code in a file called main.py. You can run python main.py -h to see a list of available options. Train a new model and save it (python main.py --save-model True) or download the pre-trained model available here. Either way, move the saved model in a new folder called artefacts.
Next, we need to write a custom handler to run the inference on your model. Copy and paste the following code in a new file called handler.py. This code runs the inference on greyscale images, like the MNIST ones:
Next, we need to create a torch model archive using the torch-model-archiver utility. This command archives the provided model (artefacts/model.pt) into a .mar file:
torch-model-archiver --model-name mnist --version 1.0 --model-file model.py --serialized-file artefacts/model.pt --handler handler.py
Create a model_store folder and move the archived model inside:
mkdir model_storemv mnist.mar model_store/
Finally, start the server and query it for answers. You can download some test data here. To use them, just put them in a new folder called test_data:
torchserve --start --model-store model_store --models mnist=mnist.marcurl http://127.0.0.1:8080/predictions/mnist -T test_data/0.png
When you are done, stop the server:
torchserve --stop
In this story, we introduced TorchServe, a flexible and easy to use tool for serving PyTorch models. We saw what it is, what it offers and how we can leverage its utilities to serve PyTorch models via a REST endpoint. Finally, we examined its latest release (version 0.1.1) and the newest supported features and put it to the test with an MNIST example. To get your hands dirty, dive in the documentation and official examples.
Learning Rate is my weekly newsletter for those who are curious about the world of AI and MLOps. You’ll hear from me every Friday with updates and thoughts on the latest AI news, research, repos and books. Subscribe here!
My name is Dimitris Poulopoulos and I’m a machine learning researcher at BigDataStack and PhD(c) at the University of Piraeus, Greece. I have worked on designing and implementing AI and software solutions for major clients such as the European Commission, Eurostat, IMF, the European Central Bank, OECD, and IKEA. If you are interested in reading more posts about Machine Learning, Deep Learning, Data Science and DataOps follow me on Medium, LinkedIn or @james2pl on twitter.