86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
|
# Copyright (c) OpenMMLab. All rights reserved.
|
||
|
|
import multiprocessing as mp
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import pytest
|
||
|
|
from mmcv import Config
|
||
|
|
|
||
|
|
from mmseg.utils import setup_multi_processes
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize('workers_per_gpu', (0, 2))
|
||
|
|
@pytest.mark.parametrize(('valid', 'env_cfg'), [(True,
|
||
|
|
dict(
|
||
|
|
mp_start_method='fork',
|
||
|
|
opencv_num_threads=0,
|
||
|
|
omp_num_threads=1,
|
||
|
|
mkl_num_threads=1)),
|
||
|
|
(False,
|
||
|
|
dict(
|
||
|
|
mp_start_method=1,
|
||
|
|
opencv_num_threads=0.1,
|
||
|
|
omp_num_threads='s',
|
||
|
|
mkl_num_threads='1'))])
|
||
|
|
def test_setup_multi_processes(workers_per_gpu, valid, env_cfg):
|
||
|
|
# temp save system setting
|
||
|
|
sys_start_mehod = mp.get_start_method(allow_none=True)
|
||
|
|
sys_cv_threads = cv2.getNumThreads()
|
||
|
|
# pop and temp save system env vars
|
||
|
|
sys_omp_threads = os.environ.pop('OMP_NUM_THREADS', default=None)
|
||
|
|
sys_mkl_threads = os.environ.pop('MKL_NUM_THREADS', default=None)
|
||
|
|
|
||
|
|
config = dict(data=dict(workers_per_gpu=workers_per_gpu))
|
||
|
|
config.update(env_cfg)
|
||
|
|
cfg = Config(config)
|
||
|
|
setup_multi_processes(cfg)
|
||
|
|
|
||
|
|
# test when cfg is valid and workers_per_gpu > 0
|
||
|
|
# setup_multi_processes will work
|
||
|
|
if valid and workers_per_gpu > 0:
|
||
|
|
# test config without setting env
|
||
|
|
|
||
|
|
assert os.getenv('OMP_NUM_THREADS') == str(env_cfg['omp_num_threads'])
|
||
|
|
assert os.getenv('MKL_NUM_THREADS') == str(env_cfg['mkl_num_threads'])
|
||
|
|
# when set to 0, the num threads will be 1
|
||
|
|
assert cv2.getNumThreads() == env_cfg[
|
||
|
|
'opencv_num_threads'] if env_cfg['opencv_num_threads'] > 0 else 1
|
||
|
|
if platform.system() != 'Windows':
|
||
|
|
assert mp.get_start_method() == env_cfg['mp_start_method']
|
||
|
|
|
||
|
|
# revert setting to avoid affecting other programs
|
||
|
|
if sys_start_mehod:
|
||
|
|
mp.set_start_method(sys_start_mehod, force=True)
|
||
|
|
cv2.setNumThreads(sys_cv_threads)
|
||
|
|
if sys_omp_threads:
|
||
|
|
os.environ['OMP_NUM_THREADS'] = sys_omp_threads
|
||
|
|
else:
|
||
|
|
os.environ.pop('OMP_NUM_THREADS')
|
||
|
|
if sys_mkl_threads:
|
||
|
|
os.environ['MKL_NUM_THREADS'] = sys_mkl_threads
|
||
|
|
else:
|
||
|
|
os.environ.pop('MKL_NUM_THREADS')
|
||
|
|
|
||
|
|
elif valid and workers_per_gpu == 0:
|
||
|
|
|
||
|
|
if platform.system() != 'Windows':
|
||
|
|
assert mp.get_start_method() == env_cfg['mp_start_method']
|
||
|
|
assert cv2.getNumThreads() == env_cfg[
|
||
|
|
'opencv_num_threads'] if env_cfg['opencv_num_threads'] > 0 else 1
|
||
|
|
assert 'OMP_NUM_THREADS' not in os.environ
|
||
|
|
assert 'MKL_NUM_THREADS' not in os.environ
|
||
|
|
if sys_start_mehod:
|
||
|
|
mp.set_start_method(sys_start_mehod, force=True)
|
||
|
|
cv2.setNumThreads(sys_cv_threads)
|
||
|
|
if sys_omp_threads:
|
||
|
|
os.environ['OMP_NUM_THREADS'] = sys_omp_threads
|
||
|
|
if sys_mkl_threads:
|
||
|
|
os.environ['MKL_NUM_THREADS'] = sys_mkl_threads
|
||
|
|
|
||
|
|
else:
|
||
|
|
assert mp.get_start_method() == sys_start_mehod
|
||
|
|
assert cv2.getNumThreads() == sys_cv_threads
|
||
|
|
assert 'OMP_NUM_THREADS' not in os.environ
|
||
|
|
assert 'MKL_NUM_THREADS' not in os.environ
|