TorchServe & Flask for Image Style Transfer

An example of web app backed by TorchServe model server

In the .

(btw, if you want to see a very simple example of web app + TorchServe for image classification then checkout to the feature/classify branch in git)

In this post I assume you’re already familiar with basics of TorchServe (handler, model files etc.). If not refer to the previous post.

Neural style transfer

In case you don’t remember how the style transfer works here is a short description. It’s important to get the high-level overview in order to understand what will be going on in the handler for TorchServe.

“Inference” in style transfer is not just one pass of the input tensor through the net. Instead a tensor (which is going to be the output picture at the end) is passed many times and the tensor itself is modified so that to minimize the content and style loss functions. At each iteration the image is changed. And the “inference” is a sequence of the iterations.

For the solution it means the following:

• the inference function in handler will be pretty complex. It’d be messy to put everything in handler.py. So, I’ll put it in additional module and show how to include it into TorchServe artifacts
• a side effect: “inference” will take some time. To not ask a user for waiting a minute while Flask will reload the entire page I’ll use ajax request from browser to the app. So, the page in browser won’t be frozen

Model file

The VGG19 pretrained model is used in the solution. Generally speaking I just followed the style transfer official example from PyTorch:

Postprocess function

The function is pretty straightforward. The only thing here is how to pass the image to the Flask app so that it can be correctly read from json. Passing it as bytes buffer worked for me.

As I mentioned to not waiting long time for reloading the whole page the generation request is sent as ajax. So, there is also a simple JQuery script for that:

Well, looks interesting 🤪

Conclusion

This post together with the previous one shows how to serve your ML models with a dedicated serving framework and how to use an approach of model server detached from application server.

It’s shown that TorchServe allows flexible customisation of pre-, postprocessing and inference functions. Thus , you can incorporate any complex logic of your models.

Also the GitHub repo with the end-to-end example can be used as a starting point for your own experiments.

As a conclusion let me mention just few benefits that an approach with a model server offers:

• more efficient use of hardware (e.g. model server can be deployed on a machine with GPUs while the application server may not need it)
• dedicated serving frameworks offer features to serve models at scale (e.g. threads and workers in TorchServe)
• serving frameworks also provide the features to speed up the development (and to not reinvent the wheel): model versioning, logs, metrics etc.
• message queuing service can be easily added to scale the solution
• dev and ML/DS teams can work more independently

It’s not the complete list but just few reasons to think about serving ML models with dedicated frameworks.

Hope you could find some helpful and practical stuff in this post.

on Medium, where people are continuing the conversation by highlighting and responding to this story.