Imposta il numero massimo di thread in fase di esecuzione su numpy/openblas

Question

Mi piacerebbe sapere se è possibile modificare in runtime (Python) il numero massimo di thread utilizzati da OpenBLAS dietro numpy?

So che è possibile impostarlo prima di eseguire l'interprete tramite la variabile di ambiente OMP_NUM_THREADS, ma mi piacerebbe cambiarlo in fase di esecuzione.

In genere, quando si utilizza MKL invece di OpenBLAS, è possibile:

import mkl 
mkl.set_num_threads(n) 

Answer 1

Si può fare questo chiamando la funzione openblas_set_num_threads utilizzando ctypes. Spesso mi trovo a voler fare questo, così ho scritto un po 'di gestore di contesto:

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) 

è possibile utilizzarlo in questo modo:

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

Come accennato nei commenti, openblas_get_num_threads e openblas_get_num_procs erano molto nuovo funzionalità al momento della scrittura, e potrebbe quindi non essere disponibile a meno che non abbiate compilato OpenBLAS dall'ultima versione del codice sorgente.