1# mypy: allow-untyped-defs 2r"""Utility classes & functions for data loading. Code in this folder is mostly used by ../dataloder.py. 3 4A lot of multiprocessing is used in data loading, which only supports running 5functions defined in global environment (py2 can't serialize static methods). 6Therefore, for code tidiness we put these functions into different files in this 7folder. 8""" 9 10import atexit 11import sys 12 13# old private location of the ExceptionWrapper that some users rely on: 14from torch._utils import ExceptionWrapper 15 16 17IS_WINDOWS = sys.platform == "win32" 18 19 20MP_STATUS_CHECK_INTERVAL = 5.0 21r"""Interval (in seconds) to check status of processes to avoid hanging in 22 multiprocessing data loading. This is mainly used in getting data from 23 another process, in which case we need to periodically check whether the 24 sender is alive to prevent hanging.""" 25 26 27python_exit_status = False 28r"""Whether Python is shutting down. This flag is guaranteed to be set before 29the Python core library resources are freed, but Python may already be exiting 30for some time when this is set. 31 32Hook to set this flag is `_set_python_exit_flag`, and is inspired by a similar 33hook in Python 3.7 multiprocessing library: 34https://github.com/python/cpython/blob/d4d60134b29290049e28df54f23493de4f1824b6/Lib/multiprocessing/util.py#L277-L327 35""" 36 37 38try: 39 import numpy 40 41 HAS_NUMPY = True 42except ModuleNotFoundError: 43 HAS_NUMPY = False 44 45 46def _set_python_exit_flag(): 47 global python_exit_status 48 python_exit_status = True 49 50 51atexit.register(_set_python_exit_flag) 52 53 54from . import collate, fetch, pin_memory, signal_handling, worker 55