Spaces:
Runtime error
Runtime error
File size: 2,668 Bytes
8aa4f1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
import platform
import warnings
import torch.multiprocessing as mp
def set_multi_processing(
mp_start_method: str = "fork", opencv_num_threads: int = 0, distributed: bool = True
) -> None:
"""Set multi-processing related environment.
This function is refered from https://github.com/open-mmlab/mmengine/blob/main/mmengine/utils/dl_utils/setup_env.py
Args:
mp_start_method (str): Set the method which should be used to start
child processes. Defaults to 'fork'.
opencv_num_threads (int): Number of threads for opencv.
Defaults to 0.
distributed (bool): True if distributed environment.
Defaults to False.
""" # noqa
# set multi-process start method as `fork` to speed up the training
if platform.system() != "Windows":
current_method = mp.get_start_method(allow_none=True)
if current_method is not None and current_method != mp_start_method:
warnings.warn(
f"Multi-processing start method `{mp_start_method}` is "
f"different from the previous setting `{current_method}`."
f"It will be force set to `{mp_start_method}`. You can "
"change this behavior by changing `mp_start_method` in "
"your config."
)
mp.set_start_method(mp_start_method, force=True)
try:
import cv2
# disable opencv multithreading to avoid system being overloaded
cv2.setNumThreads(opencv_num_threads)
except ImportError:
pass
# 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 distributed:
omp_num_threads = 1
warnings.warn(
"Setting OMP_NUM_THREADS environment variable for each process"
f" to be {omp_num_threads} in default, to avoid your system "
"being overloaded, please further tune the variable for "
"optimal 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 distributed:
mkl_num_threads = 1
warnings.warn(
"Setting MKL_NUM_THREADS environment variable for each process"
f" to be {mkl_num_threads} in default, to avoid your system "
"being overloaded, please further tune the variable for "
"optimal performance in your application as needed."
)
os.environ["MKL_NUM_THREADS"] = str(mkl_num_threads) |