-
Notifications
You must be signed in to change notification settings - Fork 452
Introduce Classification Factory and Simplify Model Imports #4456
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
Introduce Classification Factory and Simplify Model Imports #4456
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks Kirill for the great work on enhancing the classification model factories!
I added some comments related to
- Two EfficientNet classes (factory vs backbone) created confusion
- Using Literal types for model names
Something like:
MobileNetV3: Literal["mobilenetv3_small", "mobilenetv3_large"] (2 variants)
VisionTransformer: Literal["vit-tiny", "vit-small", "vit-base", "vit-large", "dinov2-small", "dinov2-base", "dinov2-large", "dinov2-giant"] (8 variants)
...
- Improve docstring to expose 50+ TorchVision models are hidden from users: I think either
list_models
and write a example docstring or just explain in the docstring where could user find those available models
Lastly, could you please add usage examples in the docstrings to demonstrate the improvements? I think this would help users understand how to use the enhanced factory APIs:
Something like:
class MobileNetV3:
"""Factory class for MobileNetV3 models.
Args:
model_name (Literal["mobilenetv3_small", "mobilenetv3_large"]): Model variant.
Examples:
>>> # Basic usage with type safety
>>> model = MobileNetV3(
... task="multi_class",
... model_name="mobilenetv3_small", # IDE shows both options
... label_info=10
... )
>>> # Multi-label classification
>>> model = MobileNetV3(
... task="multi_label",
... model_name="mobilenetv3_large",
... label_info=[1, 5, 10] # Multi-label setup
... )
"""
But also for TVModel, VisionTransformer, etc.
@eugene123tw, thank you for your review! I applied your suggestions. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks great! Thank you for your efforts! :)
Summary
How to test
Checklist
License
Feel free to contact the maintainers if that's a concern.