Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add pretrained weights for MobileNet v2 and v3 #288

Open
wants to merge 3 commits into
base: master
Choose a base branch
from

Conversation

paulnovo
Copy link

@paulnovo paulnovo commented Feb 1, 2025

This PR adds pretrained weights for MobileNet v2 and MobileNet v3 (both small and large) using the port_torchvision.jl script to generate the weights from torchvision. The PR also fixes MobileNet v3 in Metalhead to match the original paper and torchvision implementation.

Adding pretrained weights checklist

  • Train your model or port the weights from another framework.
  • Save the model state using BSON.jl with BSON.@save "modelname.bson" model_state=Flux.state(model). It is important that your model is saved under the key model_state.
  • Compress the saved model as a tarball using tar -cvzf modelname.tar.gz modelname.bson.
  • Obtain the SHAs (see the Pkg docs). Edit the Artifacts.toml file in the Metalhead.jl repository and add entry for your model. You can leave the URL empty for now.
  • Open a PR on Metalhead.jl. Be sure to ping a maintainer (e.g. @darsnack or @theabhirath) to let us know that you are adding a pre-trained weight. We will create a model repository on HuggingFace if it does not already exist.
  • Open a PR to the corresponding HuggingFace repo. Do this by going to the "Community" tab in the HuggingFace repository. PRs and discussions are shown as the same thing in the HuggingFace web app. You can use your local Git program to make clone the repo and make PRs if you wish. Check out the guide on PRs to HuggingFace for more information.
  • Copy the download URL for the model file that you added to HuggingFace. Make sure to grab the URL for a specific commit and not for the main branch.
  • Update your Metalhead.jl PR by adding the URL to the Artifacts.toml.
  • If the tests pass for your weights, we will merge your PR! Your model should pass the acctest function in the Metalhead.jl test suite. If your model already exists in the repo, then these tests are already in place, and you can add your model configuration to the PRETRAINED_MODELS list in the runtests.jl file. Please refer to the ResNet tests as an example.

@CarloLucibello
Copy link
Member

I create an huggingface repo where you can upload the weights
https://huggingface.co/FluxML/mobilenet

@paulnovo
Copy link
Author

I create an huggingface repo where you can upload the weights https://huggingface.co/FluxML/mobilenet

Thank you @CarloLucibello, I have uploaded the weights and updated the PR. I think this is good to go, but let me know if anything else is needed.

* Set epsilon to 0.001 and momentum to 0.01 in all BatchNorm layers
* MobileNet V3 uses a ReLu activation in all squeeze excitation layers
* Fix stride in final mbconv block of both small and large models, this
  should be set to 2 instead of 1
* Adds pre-trained weights for MobileNet v3 small and large models
* Adds pre-trained weights for MobileNet v2
* Adds, commented, lines to the model_list in the port_torchvision.jl
  script for generating these weights from torchvision models
* Replaces the test image (of a guitar) in the pytorch2flux.jl script,
  the link is currently unavailable
* Add models to PRETRAINED_MODELS to enable acctest
* Update pretrained status in README
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants