Our team at Flapmax is always more than happy to take on machine learning challenges. We all firmly believe (and can see) the impact artificial intelligence can and does have on the world. When we saw that one feature request on our roadmap lined up with this Hackathon, we thought it was a perfect opportunity to take part and implement PyTorch support into Measure.
What it does
Measure is our managed cloud service for benchmarking machine learning models, including computer vision, and NLP models, against a variety of hardware platforms. It makes it easy for machine learning engineers to get meaningful feedback on the performance of their models across different machines.
Now, we've made PyTorch support available. People using Measure can upload a PyTorch model of their choice and run inferences on various machines, using various dataset formats, to derive key insights that will help them make informed decisions regarding further development or use of their model.
How we built it
- Python, FastAPI, PostgreSQL, Blob, OpenVINO
The integration used some existing architecture regarding the management of models through our REST API. The fun part was getting to convert and optimize PyTorch models for use with our inference tool.
To use PyTorch models with OpenVINO we need to do two things, convert the model into ONNX format, and then use OpenVINO's model-optimizer to optimize the ONNX model for inferences.
The conversion itself can be done with a simple helper function that loads a given PyTorch model and exports it using the torch module. For the optimization, we ultimately created a microservice that takes in an ONNX model and returns an OpenVINO optimized model.
Once the optimizer microservice returns the OpenVINO optimized model, we can store it for later inference use.
Challenges we ran into
We had two main challenges during this hackathon.
Loading user's PyTorch models
How to go about loading a user's PyTorch model wasn't immediately apparent at first. It was assumed that PyTorch would have a method for handling this for us. Later, it seemed to be the case that there was no such method.
We decided to use Python's importlib module to access a given PyTorch model class by dynamically loading them into our program, which we then populated with a given state dictionary. This class could be loaded into the program by parsing the first instance of "class" in a given model file to get the class name for instantiation.
Using OpenVINO with multiple frameworks
Because OpenVINO's optimizer requires different environments for different frameworks, there was some initial confusion in the planning phase of our PyTorch integration. We had originally opted for using Python's virtual environments but, having recently dockerized our backend, the complexity of referencing system files directly was an issue.
So the obvious solution after that small bout of confusion was to create a Dockerized service specifically for ONNX optimization! (Because remember, we convert the PyTorch model into an ONNX model to optimize it with OpenVINO)
This worked out extremely well and helped to keep our backend architecture simple and easy to use.
Accomplishments that we're proud of
The integration itself only took several days which was great! And we got a really useful feature that is well implemented as a bonus.
What we learned
During this process, having come from using TensorFlow, it was interesting to see PyTorch's format for saving and loading models. Being able to call them directly as class instances in a program was strange at first but given some consideration seems highly useful and intuitive for a Python developer.
What's next for Measure
Measure's next steps are to continue growing in functionality with three key focus areas.
- Supporting more frameworks such as ONNX directly
- Quality of life features such as real-time process updates with Websockets
- More analytic metrics such as accuracy of inference when applicable
We're more than happy to add more based on any of your suggestions. And we are making Measure (Early-Access version) available to all participants of this hackathon. To get started, go to https://www.flapmax.com/platforms/measure
Thanks so much and be well!