Overview
from mfai.pytorch.models import <Model> currently appears to eagerly import every model module in the subpackage at import time. As a result, importing a lightweight model such as HalfUNet, which only needs torch and einops, can also pull in heavier dependencies such as monai, torchvision, timm, torch-geometric, etc.
This is a feature request to make model imports lazy, so users only pay the import cost for the model they actually use.
Problem Statement
For example:
python -c "import time, torch; t=time.perf_counter(); from mfai.pytorch.models import HalfUNet; print(f'{(time.perf_counter()-t):.2f} s')"Here torch is pre-imported to isolate what mfai.pytorch.models adds. The exact timing varies by system, but the import can be noticeably slow in fresh Python sessions.
The likely structural cause is the eager pkgutil.walk_packages + importlib.import_module loop in mfai/pytorch/models/__init__.py, which imports every model module to populate the registry. There are also top-level import onnx / import onnxruntime statements in mfai/pytorch/__init__.py, although these only seem to be needed by specific functions.
This is probably not a major issue during large training runs, where startup time is negligible compared with training time. However, it becomes quite noticeable in interactive workflows. For example, during an interactive FRAIM workshop using notebooks, importing from mfai.pytorch.models felt very slow compared with the actual lightweight model being used. This kind of overhead is also inconvenient in notebooks, shell loops, sweeps, and quick experiments where fresh interpreters are launched repeatedly.
Proposed Solution
Make importing mfai.pytorch.models avoid loading all model modules immediately, while keeping the usual import paths working, for example:
from mfai.pytorch.models import HalfUNet
from mfai.pytorch.models.half_unet import HalfUNet
Ideally, the registry/discovery mechanism could become lazy as well, so it does not require importing every model and all optional model-specific dependencies during package import.
Since this concerns the package import structure and registry design, the exact implementation is probably an architectural decision for the maintainers. Possible approaches could include module-level lazy attribute loading, a discover-on-first-lookup registry, or another pattern that fits the package design.
Similarly, onnx and onnxruntime could perhaps be imported only inside the functions that actually need them.
A useful longer-term consequence would be that heavy dependencies needed only by some models, such as monai, torchvision, timm, torch-geometric, or axial-attention, could potentially become optional extras rather than hard dependencies. That would allow a leaner pip install mfai and better failure isolation: a missing dependency for model X would not prevent importing model Y.
(Optional) References and Resources
Contributor Willingness
Overview
from mfai.pytorch.models import <Model>currently appears to eagerly import every model module in the subpackage at import time. As a result, importing a lightweight model such asHalfUNet, which only needstorchandeinops, can also pull in heavier dependencies such asmonai,torchvision,timm,torch-geometric, etc.This is a feature request to make model imports lazy, so users only pay the import cost for the model they actually use.
Problem Statement
For example:
python -c "import time, torch; t=time.perf_counter(); from mfai.pytorch.models import HalfUNet; print(f'{(time.perf_counter()-t):.2f} s')"Here
torchis pre-imported to isolate whatmfai.pytorch.modelsadds. The exact timing varies by system, but the import can be noticeably slow in fresh Python sessions.The likely structural cause is the eager
pkgutil.walk_packages+importlib.import_moduleloop inmfai/pytorch/models/__init__.py, which imports every model module to populate the registry. There are also top-levelimport onnx/import onnxruntimestatements inmfai/pytorch/__init__.py, although these only seem to be needed by specific functions.This is probably not a major issue during large training runs, where startup time is negligible compared with training time. However, it becomes quite noticeable in interactive workflows. For example, during an interactive FRAIM workshop using notebooks, importing from
mfai.pytorch.modelsfelt very slow compared with the actual lightweight model being used. This kind of overhead is also inconvenient in notebooks, shell loops, sweeps, and quick experiments where fresh interpreters are launched repeatedly.Proposed Solution
Make importing
mfai.pytorch.modelsavoid loading all model modules immediately, while keeping the usual import paths working, for example:Ideally, the registry/discovery mechanism could become lazy as well, so it does not require importing every model and all optional model-specific dependencies during package import.
Since this concerns the package import structure and registry design, the exact implementation is probably an architectural decision for the maintainers. Possible approaches could include module-level lazy attribute loading, a discover-on-first-lookup registry, or another pattern that fits the package design.
Similarly,
onnxandonnxruntimecould perhaps be imported only inside the functions that actually need them.A useful longer-term consequence would be that heavy dependencies needed only by some models, such as
monai,torchvision,timm,torch-geometric, oraxial-attention, could potentially become optional extras rather than hard dependencies. That would allow a leanerpip install mfaiand better failure isolation: a missing dependency for model X would not prevent importing model Y.(Optional) References and Resources
Contributor Willingness
I am happy to help test or provide more details if useful.