diff --git a/configs/_base_/default_runtime.py b/configs/_base_/default_runtime.py index 65fb632ed..3a5f0e9ef 100644 --- a/configs/_base_/default_runtime.py +++ b/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' diff --git a/mmtrack/core/utils/__init__.py b/mmtrack/core/utils/__init__.py index 274d9bf50..7e4c0e777 100644 --- a/mmtrack/core/utils/__init__.py +++ b/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' +] diff --git a/mmtrack/core/utils/misc.py b/mmtrack/core/utils/misc.py new file mode 100644 index 000000000..20c03379b --- /dev/null +++ b/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) diff --git a/tools/test.py b/tools/test.py index 3a605615f..7e177c5b4 100644 --- a/tools/test.py +++ b/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): diff --git a/tools/train.py b/tools/train.py index 29597b5a3..893537f96 100644 --- a/tools/train.py +++ b/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