Accelerating training speed.

configs/_base_/default_runtime.py

@@ -15,3 +15,8 @@
 load_from = None
 resume_from = None
 workflow = [('train', 1)]
+
+# disable opencv multithreading to avoid system being overloaded
+opencv_num_threads = 0
+# set multi-process start method as `fork` to speed up the training
+mp_start_method = 'fork'

mmtrack/core/utils/__init__.py

@@ -1,5 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .image import crop_image
+from .misc import setup_multi_processes
 from .visualization import imshow_mot_errors, imshow_tracks
 
-__all__ = ['crop_image', 'imshow_tracks', 'imshow_mot_errors']
+__all__ = [
+    'crop_image', 'imshow_tracks', 'imshow_mot_errors', 'setup_multi_processes'
+]

mmtrack/core/utils/misc.py

@@ -0,0 +1,38 @@
+import multiprocessing as mp
+import os
+import platform
+import warnings
+
+import cv2
+
+
+def setup_multi_processes(cfg):
+    # 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')
+        mp.set_start_method(mp_start_method)
+
+    # 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)

tools/test.py

@@ -13,6 +13,7 @@
                          wrap_fp16_model)
 from mmdet.apis import set_random_seed
 
+from mmtrack.core import setup_multi_processes
 from mmtrack.datasets import build_dataset
 
 
@@ -111,6 +112,9 @@ def main():
     if args.cfg_options is not None:
         cfg.merge_from_dict(args.cfg_options)
 
+    # set multi-process settings
+    setup_multi_processes(cfg)
+
     # set random seeds. Force setting fixed seed and deterministic=True in SOT
     # configs
     if cfg.get('cudnn_benchmark', False):

tools/train.py

@@ -13,6 +13,7 @@
 
 from mmtrack import __version__
 from mmtrack.apis import init_random_seed
+from mmtrack.core import setup_multi_processes
 from mmtrack.datasets import build_dataset
 from mmtrack.utils import collect_env, get_root_logger
 
@@ -83,6 +84,10 @@ def main():
         from mmtrack.models import build_model
     if args.cfg_options is not None:
         cfg.merge_from_dict(args.cfg_options)
+
+    # set multi-process settings
+    setup_multi_processes(cfg)
+
     # set cudnn_benchmark
     if cfg.get('cudnn_benchmark', False):
         torch.backends.cudnn.benchmark = True