Ограничить количество потоков в numpy

Кажется, что моя библиотека numpy использует 4 потока, а установка OMP_NUM_THREADS=1 не останавливает это.

numpy.show_config() дает мне следующие результаты:

atlas_threads_info:
    libraries = ['lapack', 'ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\"3.8.4\\""')]
    language = f77
    include_dirs = ['/usr/include']
blas_opt_info:
    libraries = ['ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\"3.8.4\\""')]
    language = c
    include_dirs = ['/usr/include']
atlas_blas_threads_info:
    libraries = ['ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\"3.8.4\\""')]
    language = c
    include_dirs = ['/usr/include']
openblas_info:
  NOT AVAILABLE
lapack_opt_info:
    libraries = ['lapack', 'ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\"3.8.4\\""')]
    language = f77
    include_dirs = ['/usr/include']

Итак, я знаю, что он использует blas, но я не могу понять, как заставить его использовать 1 поток для матричного умножения.

Ответ 1

Попробуйте установить все следующее:

export MKL_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export OMP_NUM_THREADS=1

Иногда бывает сложно увидеть, где именно вводится многопоточность.

Ответ 2

Существует более 3 упомянутых переменных среды. Ниже приведен полный список переменных среды и пакет, в котором эта переменная используется для управления количеством создаваемых потоков. Обратите внимание, что вам нужно установить эти переменные перед выполнением import numpy:

OMP_NUM_THREADS: openmp,
OPENBLAS_NUM_THREADS: openblas,
MKL_NUM_THREADS: mkl,
VECLIB_MAXIMUM_THREADS: accelerate,
NUMEXPR_NUM_THREADS: numexpr

Итак, на практике вы можете сделать:

import os
os.environ["OMP_NUM_THREADS"] = "4" # export OMP_NUM_THREADS=4
os.environ["OPENBLAS_NUM_THREADS"] = "4" # export OPENBLAS_NUM_THREADS=4 
os.environ["MKL_NUM_THREADS"] = "6" # export MKL_NUM_THREADS=6
os.environ["VECLIB_MAXIMUM_THREADS"] = "4" # export VECLIB_MAXIMUM_THREADS=4
os.environ["NUMEXPR_NUM_THREADS"] = "6" # export NUMEXPR_NUM_THREADS=6

Обратите внимание, что с ноября 2018 года разработчики Numpy работают над тем, чтобы сделать это возможным и после того, как вы сделаете import numpy. Я обновлю этот пост, как только они внесут эти изменения.

Ответ 3

Что касается этого в скрипте Python, а не в приглашении bash, для этого потока вы можете сделать следующее (те же команды, что и в ответе выше):

import os
os.environ["MKL_NUM_THREADS"] = "1" 
os.environ["NUMEXPR_NUM_THREADS"] = "1" 
os.environ["OMP_NUM_THREADS"] = "1"

но вы должны поставить это, прежде чем import numpy. Видимо, numpy только проверяет это при импорте.

(это сообщение отправлено как ответ на основании комментария @kηives выше.)

Ответ 4

Я смог исправить это во время выполнения следующим образом:

import mkl
mkl.set_num_threads(1)

Я использую следующий код, чтобы уменьшить вероятность появления этого фрагмента в сценариях/пакетах:

try:
    import mkl
    mkl.set_num_threads(1)
except:
    pass