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

feat: add lazy imports for faster startup time #160

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ einops
faster-coco-eval
graphviz
hydra-core
lazy-loader
lightning
loguru
numpy
Expand Down
109 changes: 77 additions & 32 deletions yolo/__init__.py
Original file line number Diff line number Diff line change
@@ -1,33 +1,78 @@
from yolo.config.config import Config, NMSConfig
from yolo.model.yolo import create_model
from yolo.tools.data_loader import AugmentationComposer, create_dataloader
from yolo.tools.drawer import draw_bboxes
from yolo.tools.solver import TrainModel
from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter
from yolo.utils.deploy_utils import FastModelLoader
from yolo.utils.logging_utils import (
ImageLogger,
YOLORichModelSummary,
YOLORichProgressBar,
"""
The MIT YOLO rewrite
"""

__autogen__ = """
mkinit ~/code/YOLO-v9/yolo/__init__.py --nomods --write --lazy-loader

# Check to see how long it takes to run a simple help command
time python -m yolo.lazy --help
"""

__submodules__ = {
'config.config': ['Config', 'NMSConfig'],
'model.yolo': ['create_model'],
'tools.data_loader': ['AugmentationComposer', 'create_dataloader'],
'tools.drawer': ['draw_bboxes'],
'tools.solver': ['TrainModel'],
'utils.bounding_box_utils': ['Anc2Box', 'Vec2Box', 'bbox_nms', 'create_converter'],
'utils.deploy_utils': ['FastModelLoader'],
'utils.logging_utils': [
'ImageLogger', 'YOLORichModelSummary',
'YOLORichProgressBar',
'validate_log_directory'
],
'utils.model_utils': ['PostProcess'],
}


import lazy_loader

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://peps.python.org/pep-0008/#imports

Imports are always put at the top of the file, just after any module comments and docstrings, and before module globals and constants.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this case it is done due a quirk of how mkinit autogenerates code, so if you want the lazy-loader dependency path I can fix that. Otherwise, that would be replaced with a definition of the function that does the same thing. Technically none of the constants are needed, they are just convenient to automatically update the __init__.py with mkinit.

However, it's important to note: always adhering to this rule can have significant startup-time consequences. I moved some imports after the call to setup inside of the main function here: https://github.com/MultimediaTechLab/YOLO/pull/160/files#diff-029d4708225fe34b2045fd6a94297034336187c511ca506dccd687f5b6fca3c3R16 and this is important for a faster startup time.



__getattr__, __dir__, __all__ = lazy_loader.attach(
__name__,
submodules={},
submod_attrs={
'config.config': [
'Config',
'NMSConfig',
],
'model.yolo': [
'create_model',
],
'tools.data_loader': [
'AugmentationComposer',
'create_dataloader',
],
'tools.drawer': [
'draw_bboxes',
],
'tools.solver': [
'TrainModel',
],
'utils.bounding_box_utils': [
'Anc2Box',
'Vec2Box',
'bbox_nms',
'create_converter',
],
'utils.deploy_utils': [
'FastModelLoader',
],
'utils.logging_utils': [
'ImageLogger',
'YOLORichModelSummary',
'YOLORichProgressBar',
'validate_log_directory',
],
'utils.model_utils': [
'PostProcess',
],
},
)
from yolo.utils.model_utils import PostProcess

all = [
"create_model",
"Config",
"YOLORichProgressBar",
"NMSConfig",
"YOLORichModelSummary",
"validate_log_directory",
"draw_bboxes",
"Vec2Box",
"Anc2Box",
"bbox_nms",
"create_converter",
"AugmentationComposer",
"ImageLogger",
"create_dataloader",
"FastModelLoader",
"TrainModel",
"PostProcess",
]

__all__ = ['Anc2Box', 'AugmentationComposer', 'Config', 'FastModelLoader',
'ImageLogger', 'NMSConfig', 'PostProcess', 'TrainModel', 'Vec2Box',
'YOLORichModelSummary', 'YOLORichProgressBar', 'bbox_nms',
'create_converter', 'create_dataloader', 'create_model',
'draw_bboxes', 'validate_log_directory']
11 changes: 5 additions & 6 deletions yolo/lazy.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,19 @@
from pathlib import Path

import hydra
from lightning import Trainer
from omegaconf.dictconfig import DictConfig

project_root = Path(__file__).resolve().parent.parent
sys.path.append(str(project_root))

from yolo.config.config import Config
from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel
from yolo.utils.logging_utils import setup


@hydra.main(config_path="config", config_name="config", version_base=None)
def main(cfg: Config):
def main(cfg: DictConfig):
from yolo.utils.logging_utils import setup
callbacks, loggers, save_path = setup(cfg)

from lightning import Trainer
from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel
trainer = Trainer(
accelerator="auto",
max_epochs=getattr(cfg.task, "epoch", None),
Expand Down
Empty file added yolo/model/__init__.py
Empty file.