Ustaw maksymalną liczbę wątków w czasie wykonywania na numpy/openblas

Question

Chciałbym wiedzieć, czy w środowisku wykonawczym (Python) możliwa jest maksymalna liczba wątków używanych przez OpenBLAS za numpy?

Wiem, że możliwe jest ustawienie go przed uruchomieniem interpretera za pomocą zmiennej środowiskowej OMP_NUM_THREADS, ale chcę go zmienić w czasie wykonywania.

Zazwyczaj podczas korzystania MKL zamiast OpenBLAS, możliwe jest:

import mkl 
mkl.set_num_threads(n) 

Answer 1

Można to zrobić poprzez wywołanie funkcji openblas_set_num_threads korzystając ctypes. I często znajduję się chce to zrobić, więc napisałem trochę menedżera kontekstowego:

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) 

można go używać tak:

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

Jak wspomniano w komentarzach, openblas_get_num_threads i openblas_get_num_procs były bardzo nowy funkcje w czasie pisania, a zatem mogą nie być dostępne, chyba że skompilowałeś OpenBLAS z najnowszej wersji kodu źródłowego.