python - numpy/openblas で実行時にスレッドの最大数を設定する

Question

(Python) 実行時に、numpy の背後にある OpenBLAS によって使用されるスレッドの最大数を変更できるかどうかを知りたいですか?

環境変数を介してインタープリターを実行する前に設定できることはわかっていますOMP_NUM_THREADSが、実行時に変更したいと思います。

通常、OpenBLAS の代わりに MKL を使用すると、次のことが可能になります。

import mkl
mkl.set_num_threads(n)

Answer 1

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 をコンパイルしない限り、利用できない可能性があります。