mmpose-0.29.0/tests/test_utils.py

# Copyright (c) OpenMMLab. All rights reserved.
import multiprocessing as mp
import os
import platform
import time

import cv2
import mmcv
import torch
import torchvision
from mmcv import Config

import mmpose
from mmpose.utils import StopWatch, collect_env, setup_multi_processes


def test_collect_env():
 env_info = collect_env()
 assert env_info['PyTorch'] == torch.__version__
 assert env_info['TorchVision'] == torchvision.__version__
 assert env_info['OpenCV'] == cv2.__version__
 assert env_info['MMCV'] == mmcv.__version__
 assert '+' in env_info['MMPose']
 assert mmpose.__version__ in env_info['MMPose']


def test_stopwatch():
 window_size = 5
 test_loop = 10
 outer_time = 100
 inner_time = 100

 stop_watch = StopWatch(window=window_size)
 for _ in range(test_loop):
 with stop_watch.timeit():
 time.sleep(outer_time / 1000.)
 with stop_watch.timeit('inner'):
 time.sleep(inner_time / 1000.)

 _ = stop_watch.report()
 _ = stop_watch.report_strings()


def test_setup_multi_processes():
 # 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)

 # test config without setting env
 config = dict(data=dict(workers_per_gpu=2))
 cfg = Config(config)
 setup_multi_processes(cfg)
 assert os.getenv('OMP_NUM_THREADS') == '1'
 assert os.getenv('MKL_NUM_THREADS') == '1'
 # when set to 0, the num threads will be 1
 assert cv2.getNumThreads() == 1
 if platform.system() != 'Windows':
 assert mp.get_start_method() == 'fork'

 # test num workers <= 1
 os.environ.pop('OMP_NUM_THREADS')
 os.environ.pop('MKL_NUM_THREADS')
 config = dict(data=dict(workers_per_gpu=0))
 cfg = Config(config)
 setup_multi_processes(cfg)
 assert 'OMP_NUM_THREADS' not in os.environ
 assert 'MKL_NUM_THREADS' not in os.environ

 # test manually set env var
 os.environ['OMP_NUM_THREADS'] = '4'
 config = dict(data=dict(workers_per_gpu=2))
 cfg = Config(config)
 setup_multi_processes(cfg)
 assert os.getenv('OMP_NUM_THREADS') == '4'

 # test manually set opencv threads and mp start method
 config = dict(
 data=dict(workers_per_gpu=2),
 opencv_num_threads=4,
 mp_start_method='spawn')
 cfg = Config(config)
 setup_multi_processes(cfg)
 assert cv2.getNumThreads() == 4
 assert mp.get_start_method() == 'spawn'

 # 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')