Set max number of threads at runtime on numpy/openblas

前端 未结 2 844
慢半拍i
慢半拍i 2020-12-30 07:57

I\'d like to know if it\'s possible to change at (Python) runtime the maximum number of threads used by OpenBLAS behind numpy?

I know it\'s possible to set it before

2条回答
  •  -上瘾入骨i
    2020-12-30 08:27

    We recently developed threadpoolctl, a cross platform package to do control the number of threads used in calls to C-level thread-pools in python. It works similarly to the answer by @ali_m but detects automatically the libraries that needs to be limited by looping through all loaded libraries. It also comes with introspection APIs.

    This package can be installed using pip install threadpoolctl and come with a context manager that allow you to control the number of threads used by packages such as 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
    

    you can also have finer control on different threadpools (such as differenciating blas from openmp calls).

    Note: this package is still in development and any feedback is welcomed.

提交回复
热议问题