Shortcuts

Source code for mmengine.utils.dl_utils.setup_env

# Copyright (c) OpenMMLab. All rights reserved.
import os
import platform
import warnings

import torch.multiprocessing as mp


[docs]def set_multi_processing(mp_start_method: str = 'fork', opencv_num_threads: int = 0, distributed: bool = False) -> None: """Set multi-processing related environment. Args: mp_start_method (str): Set the method which should be used to start child processes. Defaults to 'fork'. opencv_num_threads (int): Number of threads for opencv. Defaults to 0. distributed (bool): True if distributed environment. Defaults to False. """ # set multi-process start method as `fork` to speed up the training if platform.system() != 'Windows': current_method = mp.get_start_method(allow_none=True) if (current_method is not None and current_method != mp_start_method): warnings.warn( f'Multi-processing start method `{mp_start_method}` is ' f'different from the previous setting `{current_method}`.' f'It will be force set to `{mp_start_method}`. You can ' 'change this behavior by changing `mp_start_method` in ' 'your config.') mp.set_start_method(mp_start_method, force=True) try: import cv2 # disable opencv multithreading to avoid system being overloaded cv2.setNumThreads(opencv_num_threads) except ImportError: pass # setup OMP threads # This code is referred from https://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py # noqa if 'OMP_NUM_THREADS' not in os.environ and distributed: omp_num_threads = 1 warnings.warn( 'Setting OMP_NUM_THREADS environment variable for each process' f' to be {omp_num_threads} in default, to avoid your system ' 'being overloaded, please further tune the variable for ' 'optimal performance in your application as needed.') os.environ['OMP_NUM_THREADS'] = str(omp_num_threads) # setup MKL threads if 'MKL_NUM_THREADS' not in os.environ and distributed: mkl_num_threads = 1 warnings.warn( 'Setting MKL_NUM_THREADS environment variable for each process' f' to be {mkl_num_threads} in default, to avoid your system ' 'being overloaded, please further tune the variable for ' 'optimal performance in your application as needed.') os.environ['MKL_NUM_THREADS'] = str(mkl_num_threads)

© Copyright 2022, mmengine contributors. Revision 39ed23fa.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
v0.10.3
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.