Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save EricCousineau-TRI/8a2d1550f5fa4be4fed87d55a522dbf2 to your computer and use it in GitHub Desktop.
Save EricCousineau-TRI/8a2d1550f5fa4be4fed87d55a522dbf2 to your computer and use it in GitHub Desktop.
"""
Import this first, before *anything* else, if you are using `multiprocessing`
directly or indirectly and you encounter "freezing" when any threaded code uses
things like NumPy, OpenCV, etc.
This should only be be necessary when using the "fork" start method; other methods
should (hopefully) work fine without it:
https://docs.python.org/3.6/library/multiprocessing.html#contexts-and-start-methods
For more information, see:
https://stackoverflow.com/questions/17053671/python-how-do-you-stop-numpy-from-multithreading # noqa
For clarity, this should *only* be used by the *main* module! (see noisier but more robust
example below)
"""
import os
os.environ.update(
OMP_NUM_THREADS = '1',
OPENBLAS_NUM_THREADS = '1',
NUMEXPR_NUM_THREADS = '1',
MKL_NUM_THREADS = '1',
)
import numpy as np
import cv2
cv2.setNumThreads(0)
"""
Same as above, but has assertions to try and ensure that you only call this
in a `__main__` file.
Otherwise, you may "infect" programs that don't need it! (e.g. you have a program
that *doesn't* use `mp.get_context(method="fork")`, but now NumPy, OpenCV, etc.
will go super slow because you've constrained the number of threads)
Only tested on CPython 3.6.x on Ubuntu 18.04.
"""
import importlib
import inspect
import os
import sys
_IGNORE_CALLING_FILENAMES = {
# Skip importlib modules, as they're used to parse and execute the code.
'<frozen importlib._bootstrap_external>',
'<frozen importlib._bootstrap>',
importlib.__file__,
}
def _get_calling_module_name():
calling_stack = inspect.stack()[3:]
for frame_info in calling_stack:
if frame_info.filename not in _IGNORE_CALLING_FILENAMES:
break
else:
assert False, "Stack too shallow?"
m = inspect.getmodule(frame_info.frame)
return m.__name__
def _assert_main_and_import_order():
calling_name = _get_calling_module_name()
assert calling_name == "__main__", (
f"This should only be imported when the calling module is '__main__', "
f"but it is '{calling_name}'"
)
this_should_be_imported_before = {"cv2", "numpy"}
toplevel_modules = set([name for name in sys.modules if "." not in name])
already_imported = this_should_be_imported_before & toplevel_modules
assert len(already_imported) == 0, (
f"Modules should be imported *after* '{__name__}': {already_imported}"
)
def _set_env():
# This implements the actual workaround.
os.environ.update(
OMP_NUM_THREADS='1',
OPENBLAS_NUM_THREADS='1',
NUMEXPR_NUM_THREADS='1',
MKL_NUM_THREADS='1',
)
import numpy as np # noqa
import cv2 # noqa
cv2.setNumThreads(0)
_assert_main_and_import_order()
_set_env()
@EricCousineau-TRI
Copy link
Author

Welcome! FWIW I've updated this with some extra docs and a more robust "import check".

@rwarrier243
Copy link

rwarrier243 commented Jun 5, 2021 via email

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment