用 os.environ.setdefault 将原生库线程数(OMP_NUM_THREADS、OPENBLAS_NUM_THREADS、MKL_NUM_THREADS、NUMEXPR_NUM_THREADS)的默认值设置为 "1"。这样可以确保导入的代码默认不会过度占用 CPU 线程(以及巨量内存)
15 lines
303 B
Python
15 lines
303 B
Python
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
|
|
_NATIVE_THREAD_ENV_DEFAULTS = {
|
|
"OMP_NUM_THREADS": "1",
|
|
"OPENBLAS_NUM_THREADS": "1",
|
|
"MKL_NUM_THREADS": "1",
|
|
"NUMEXPR_NUM_THREADS": "1",
|
|
}
|
|
|
|
for _name, _value in _NATIVE_THREAD_ENV_DEFAULTS.items():
|
|
os.environ.setdefault(_name, _value)
|