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

Make the torchvision model architecture selectable by env var #3

Merged
merged 2 commits into from
May 14, 2024

Conversation

ethanjli
Copy link
Member

@ethanjli ethanjli commented May 13, 2024

This PR makes it possible to select the Torchvision model architecture used for model weights loaded from the models directory, by setting a new environment variable named TORCHVISION_MODEL_TYPE. If the environment variable is not set, then the efficientnet_v2_m architecture is assumed; currently, the only other supported model type is efficientnet_v2_s, which is the architecture of https://github.com/PlanktoScope/streamlit-classification-app/releases/download/models%2Fdemo-1/effv2s_no_norm_DA+sh_20patience_256x256_50ep_loss.pth referenced in README.md.

This PR also adds a feature flag to the Forklift package for the streamlit demo app to enable switching TORCHVISION_MODEL_TYPE to efficientnet_v2_s away from the default. Because the effv2s_no_norm_DA+sh_20patience_256x256_50ep_loss.pth model weights file is the only publicly available file, the demo pallet in this repo enables that feature flag by default.

@ethanjli ethanjli marked this pull request as ready for review May 13, 2024 20:36
@ethanjli
Copy link
Member Author

Here's what this PR looks like when I test it on my computer:
image

@ethanjli ethanjli merged commit 4432ff1 into main May 14, 2024
3 checks passed
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.

1 participant