pytorch / vision

Datasets, Transforms and Models specific to Computer Vision
https://pytorch.org/vision
BSD 3-Clause "New" or "Revised" License
15.98k stars 6.92k forks source link

Improve empty import time of torchvision #8497

Open bschindler opened 2 months ago

bschindler commented 2 months ago

🚀 The feature

When importing torchvision, a number of libraries are imported by default for more niche functionality of the library. To improve import time, I would favor delaying those imports to when they are needed

Motivation, pitch

In my case, it is the av library in particular that contributes to the import time:

image

(this assumes that torch, dynamo and onnx are already imported).

The import of av can easily be avoided as it is not needed by default.

Alternatives

No response

Additional context

I checked the code and I found this code here:

try:
    import av

    av.logging.set_level(av.logging.ERROR)
    if not hasattr(av.video.frame.VideoFrame, "pict_type"):
        av = ImportError(
            """\
Your version of PyAV is too old for the necessary video operations in torchvision.
If you are on Python 3.5, you will have to build from source (the conda-forge
packages are not up-to-date).  See
https://github.com/mikeboers/PyAV#installation for instructions on how to
install PyAV on your system.
"""
        )
except ImportError:
    av = ImportError(
        """\
PyAV is not installed, and is necessary for the video operations in torchvision.
See https://github.com/mikeboers/PyAV#installation for instructions on how to
install PyAV on your system.
"""
    )

The pict_type got added somewhere in the 0.5 range (released around 2020), 6.0 followed shortly. So I would suggest to change this test to not import av but the use importlib to check the version which would make this go away. This applies both to torchvision/io/video_reader.py as well as torchvision/io/video.py. I also wonder whether the logging call is still required given so much has changed since this code was written.

NicolasHug commented 2 months ago

Hi @bschindler , thanks for the report. Happy to consider a PR to import av only when the io module is imported. Or... to only import io when io is needed. Or both...

Considering the cluncky way by which we load the extensions and make them available, I suspect this is the kind of work that sounds a lot easier than it actually is, but we'll see.

bschindler commented 2 months ago

I briefly looked at it, a major issue is this line in the init file of torchvision:

from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils

This unconditionally imports io, making it part of the public interface. I could remove it from there but that would be a breaking change. Would you be in favor of that?

NicolasHug commented 1 month ago

Sorry for the late reply @bschindler . Unfortunately this can potentially break existing imports as you mentioned, so it's best not to go for that. If there is a way to import av only when it's needed, then that should be non-intrusive while still addressing the majority of the import time according to your benchmarks.