Установите максимальное количество потоков во время выполнения на numpy/openblas

Вы можете сделать это, вызвав функцию openblas_set_num_threads с помощью ctypes. Мне часто хочется это сделать, поэтому я написал небольшой контекстный менеджер:

import contextlib
import ctypes
from ctypes.util import find_library

# Prioritize hand-compiled OpenBLAS library over version in /usr/lib/
# from Ubuntu repos
try_paths = ['/opt/OpenBLAS/lib/libopenblas.so',
             '/lib/libopenblas.so',
             '/usr/lib/libopenblas.so.0',
             find_library('openblas')]
openblas_lib = None
for libpath in try_paths:
    try:
        openblas_lib = ctypes.cdll.LoadLibrary(libpath)
        break
    except OSError:
        continue
if openblas_lib is None:
    raise EnvironmentError('Could not locate an OpenBLAS shared library', 2)


def set_num_threads(n):
    """Set the current number of threads used by the OpenBLAS server."""
    openblas_lib.openblas_set_num_threads(int(n))


# At the time of writing these symbols were very new:
# https://github.com/xianyi/OpenBLAS/commit/65a847c
try:
    openblas_lib.openblas_get_num_threads()
    def get_num_threads():
        """Get the current number of threads used by the OpenBLAS server."""
        return openblas_lib.openblas_get_num_threads()
except AttributeError:
    def get_num_threads():
        """Dummy function (symbol not present in %s), returns -1."""
        return -1
    pass

try:
    openblas_lib.openblas_get_num_procs()
    def get_num_procs():
        """Get the total number of physical processors"""
        return openblas_lib.openblas_get_num_procs()
except AttributeError:
    def get_num_procs():
        """Dummy function (symbol not present), returns -1."""
        return -1
    pass


@contextlib.contextmanager
def num_threads(n):
    """Temporarily changes the number of OpenBLAS threads.

    Example usage:

        print("Before: {}".format(get_num_threads()))
        with num_threads(n):
            print("In thread context: {}".format(get_num_threads()))
        print("After: {}".format(get_num_threads()))
    """
    old_n = get_num_threads()
    set_num_threads(n)
    try:
        yield
    finally:
        set_num_threads(old_n)

Вы можете использовать его следующим образом:

with num_threads(8):
    np.dot(x, y)

Как упоминалось в комментариях, openblas_get_num_threads и openblas_get_num_procs были очень новыми функциями на момент написания, и поэтому могут быть недоступны, если вы не скомпилировали OpenBLAS из последней версии исходного кода.

Недавно мы разработали threadpoolctl, межплатформенный пакет для управления количеством потоков, используемых при вызовах уровня C. пулы потоков в python. Он работает аналогично ответу @ali_m, но автоматически определяет библиотеки, которые необходимо ограничить, перебирая все загруженные библиотеки. Он также поставляется с API-интерфейсами самоанализа.

Этот пакет можно установить с помощью pip install threadpoolctl, и он поставляется с менеджером контекста, который позволяет вам контролировать количество потоков, используемых такими пакетами, как numpy:

from threadpoolctl import threadpool_limits
import numpy as np


with threadpool_limits(limits=1, user_api='blas'):
    # In this block, calls to blas implementation (like openblas or MKL)
    # will be limited to use only one thread. They can thus be used jointly
    # with thread-parallelism.
    a = np.random.randn(1000, 1000)
    a_squared = a @ a

вы также можете лучше контролировать различные пулы потоков (например, различать вызовы blas и openmp).

Примечание. этот пакет все еще находится в разработке, и мы приветствуем любые отзывы.