Switch to side-by-side view

--- a
+++ b/ViTPose/tests/test_utils.py
@@ -0,0 +1,100 @@
+# 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')