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

Question

Я хотел бы узнать, можно ли изменить на (Python) время выполнения максимальное количество потоков, используемых OpenBLAS за numpy?

Я знаю, что его можно установить перед запуском интерпретатора через переменную окружения OMP_NUM_THREADS, но я бы хотел изменить ее во время выполнения.

Как правило, при использовании МКЛ вместо OpenBLAS, можно:

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 из последней версии исходного кода.