Sharing Axon Models on Huggingface

Finding a place to share our models

The Axon, https://github.com/elixir-nx/axon, community was looking for techniques to share models. What if we used Huggingface’s Hub as a model repository? I knew that the Hub is designed to be model framework agnostic. How agnostic are they really?

Huggingface Hub is kind of like GitHub, but focused on machine learning concerns like models and shared datasets. Their Hub model repository is built around the concept of a model card. The card contains documentation about the model, its intended uses, limitations, along with other descriptive information. The objective of the card is to help Hub users identify whether the model will fit their needs. The Hub might work for us.

Now we need to find a model to publish. There are a wealth of already developed models in a variety of frameworks. The axon_onnx package supports the ability to read in ONNX models and convert them to Axon models. Axon_onnx is a young library and doesn’t handle all ONNX capabilities. It does provide support for ResNet models.

Where can we find reusable ONNX ResNet models that we can convert to Axon? How about the ONNX model repository on Github, https://github.com/onnx/models. I focused on the ImageNet trained classification models. ResNet ImageNet models are often used as the backbone for transfer training a model in a new business domain. Having access to Axon ResNet models can help the community grow to supporting useful ML models.

We have a potential model repository and an ONNX model. I wrote some prototype code to convert onnx models into Axon models. There were some challenges. I tried a ResNet34-v2 model first and found some shape related errors. The small to medium ResNet-v1 model worked well though. At the time, I couldn’t get a ResNet101 model to convert to Axon. So, we have 3 models that we can try publishing on Huggingface.

  1. ResNet18-v1
  2. ResNet34-v1
  3. ResNet50-v1

I worked with my contact at Huggingface to make sure that they were ok with the Axon community using their Hub as an Axon model repository. There are some expectations they have. Each page has only one model. So the three models I’ve converted each need their own model page. The model card should be robust and useful. Beyond those expectations, the Axon community is welcome to use Huggingface Hub as a model repository. THANKS Huggingface!

Now let’s upload our first model. ResNet34 is my favorite. It can be pretty accurate while small enough to be relatively efficient running in a fairly intense production environment. The instructions at https://huggingface.co/docs/hub/adding-a-model worked pretty well. If, by some chance, the upload connection times out, be sure to check your email for notification of upload complete.

Huggingface has some expectations for model sharing. As their documentation states, the model card “is arguably as important as the model”. We want to help Hub users find our Axon models. The Huggingface documentation for a model card is found at: Building a Model Card. For the ResNet models, I copied the model information from the ONNX Classifcation ResNet page. Be sure to add the meta information to the README.md

---
license: apache-2.0
tags:
- Axon
- Elixir
datasets:
- ImageNet
---

We’ve updated our model card. All is good. How do we share with other Axon users? There is kind of a trick to finding Axon models. Axon isn’t big enough for Huggingface to add it to the list of libraries. So there are no buttons we can push on Huggingface’s Hub. The web page search capability finds by the model owner name or model name. We have to find Axon models by their tags. It’s possible, but cryptic, to find Axon models by tags on the Huggingface website.

https://huggingface.co/models?other=Axon
or
https://huggingface.co/models?other=Elixir

Stefano Falsetto to the Rescue!!. When I first mentioned the potential of Huggingface Hub on the Erlang Ecosystem Foundation’s slack channel. There were some concerns about our Axon users needing python in order to interact with the Hub in a command line manner. Stefano took on the challenge. He quickly worked up a prototype conversion of https://github.com/huggingface/huggingface_hub.
He even got the attention of some folks at Huggingface’s twitter account. I know my contact called me about his library. The initial version was just a shell. However, about the time that I was preparing this blog post, he’s made sufficient progress with his library, https://github.com/stefkohub/Elixir_Huggingface_hub.
The big news is that you can search for Axon models on Huggingface directly.

# List all Axon models.
iex> HuggingfaceHub.list_axon_models()

There is also a way to download a model using his library. Check out his library and consider contributing and improving Elixir’s Huggingface Hub interface.

I’ve created an Axon Organization on Huggingface. My expectation for this organization is that we can store Axon models for the more famous modeling approaches in the industry. Essentially, someday the vast majority of the ONNX model repository could be mirrored in the Axon Huggingface organization. My hope is that every model in the Axon organization has an excellent model card.

If you want to contribute to axon_onnx and the Huggingface Axon organization, try seeing which models can be easily converted to Axon. When you are successful, request to be added to the Axon organization. The admins will hopefully ask about what model you want to contribute and discuss model card expectations. Once granted access you can add your model to the organization. When you find problematic conversions, consider finding the solution and contribute the capability to axon_onnx.

Not every model belongs in the Axon organization. For example, I have plans for a pet breeds model based upon the Oxford pets dataset. Such a model wouldn’t be appropriate for the Axon organization. I’ll upload it to my Huggingface account. You can always add your model to your account.

Summary: We have Huggingface’s permission to store Axon models in their public repository.
We have Stefano’s HuggingfaceHub Elixir library to help access and interact with the Hub. We know a good source of ONNX libraries and a path to convert the libraries to Axon. What will you share?

axon, huggingface, imagenet