CODEWITHSUNDEEP
pythonpython-3.xpytorch
Maksym Titov
Maksym Titov

Reputation: 45

How to import module only when I need to use it?

I have a project which has optional dependencies, depending on the use case it may need torch library. PyTorch takes a lot of memory, it's not optimal to install it when you don't need to use it. In my project, there are parts of the code which use torch, but most of the logic does not need it. I'm trying to implement optional import so that when I need to use torch I use it, here is my idea:

def _lazy_import_torch():
    import torch

def main():
    _lazy_import_torch()
    print(f'You have torch: {torch.__version__}')

if __name__ == '__main__':
    main()

When I run this code it appears that torch was not imported:

$ python main.py
Traceback (most recent call last):
  File "main.py", line 9, in <module>
    main()
  File "main.py", line 6, in main
    print(f'You have torch: {torch.__version__}')
NameError: name 'torch' is not defined

torch is installed in my environment.

I was wondering why it did not work and how can I make it works, please help me.

Upvotes: 1

Views: 528

Answers (2)

Kemp
Kemp

Reputation: 3649

The way you're doing it, torch will be added to sys.modules but won't be visible using the name torch outside of the function where you imported it. You could look it up in the other functions using sys.modules['torch'].

If it's only used in one or two functions then you could import it at the top of those functions. You only pay the import cost once as after that it's just pulled from sys.modules. Doing that in a lot of functions would start to look very messy though.

Upvotes: 0

ForceBru
ForceBru

Reputation: 44858

This code makes torch a local name of the _lazy_import_torch function:

def _lazy_import_torch():
    import torch

After it's executed, this name is gone, so you won't be able to access it anywhere outside this function.

You could return torch from the function and then pass it around everywhere:

def _lazy_import_torch():
    import torch
    return torch

def main():
    torch = _lazy_import_torch()
    print(f'You have torch: {torch.__version__}')

Or use a global variable:

torch = None

def _lazy_import_torch():
    global torch
    import torch

If the module hasn't been imported yet, you'll get an error every time you try to use it, but after you run _lazy_import_torch, the global name torch will refer to the torch module.

Here's an example with the CSV module from the standard library:

>>> csv = None
>>> def import_csv():
...   global csv
...   import csv
... 
>>> csv
>>> import_csv()
>>> csv
<module 'csv' from '.../lib/python3.9/csv.py'>
>>>

Upvotes: 2

Related Questions