Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import os | |
| import platform | |
| import warnings | |
| import cv2 | |
| import torch.multiprocessing as mp | |
| def setup_multi_processes(cfg): | |
| """Setup multi-processing environment variables.""" | |
| # set multi-process start method as `fork` to speed up the training | |
| if platform.system() != 'Windows': | |
| mp_start_method = cfg.get('mp_start_method', 'fork') | |
| 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 ' | |
| f'this behavior by changing `mp_start_method` in your config.') | |
| mp.set_start_method(mp_start_method, force=True) | |
| # disable opencv multithreading to avoid system being overloaded | |
| opencv_num_threads = cfg.get('opencv_num_threads', 0) | |
| cv2.setNumThreads(opencv_num_threads) | |
| # 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 cfg.data.workers_per_gpu > 1: | |
| omp_num_threads = 1 | |
| warnings.warn( | |
| f'Setting OMP_NUM_THREADS environment variable for each process ' | |
| f'to be {omp_num_threads} in default, to avoid your system being ' | |
| f'overloaded, please further tune the variable for optimal ' | |
| f'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 cfg.data.workers_per_gpu > 1: | |
| mkl_num_threads = 1 | |
| warnings.warn( | |
| f'Setting MKL_NUM_THREADS environment variable for each process ' | |
| f'to be {mkl_num_threads} in default, to avoid your system being ' | |
| f'overloaded, please further tune the variable for optimal ' | |
| f'performance in your application as needed.') | |
| os.environ['MKL_NUM_THREADS'] = str(mkl_num_threads) | |