如何限制 python 庫中使用的 CPU 內核數量？

Question

我正在為 TensorFlow model（自己的代碼，而不是 tf.Data）運行 DataPipeline，使用多處理庫進行可調整的並行計算量。 盡管限制了並行進程的數量 (#CPU)，但我注意到 CPU 負載非常高（顯着高於 100%*#CPU）。

我將問題追溯到我用來計算音頻特征的 librosa 庫。

示例代碼：

import librosa
while True:
    librosa.feature.mfcc(y=np.random.normal(0, 1, (26783, )), sr=16000, n_mfcc=24, win_length=640, hop_length=640, n_fft=640)
    time.sleep(0.1)

上面簡單的最小示例代碼片段將使用我機器上大約 2000% - 3000% 的單個 CPU，如 htop 所示，這是因為 librosa 似乎在內部並行化了該操作。 由於高 CPU 負載會影響我的同事並控制我消耗的 CPU 負載，因此我想強制計算在單個（或至少可控數量的）內核上運行。

我想像：

limit_cpu_cores_for_this_process(os.getpid())
librosa.feature.mfcc(y=np.random.normal(0, 1, (26783, )), sr=16000, n_mfcc=24, win_length=640, hop_length=640, n_fft=640)

在循環內部調用（在最小示例中）或我的子進程應該只使用一個 CPU 內核，但我無法找到關於那個非常具體的主題的東西。 最后，每個限制 CPU 使用上述代碼的解決方案（在某種程度上放慢速度是可以的）都會有所幫助！

由於 librosa 在內部執行此操作，我無法直接限制 CPU 核心數，但我必須限制代碼，因為其他代碼已經並行化了。

---

編輯：自己找到了解決方案，請參閱我的答案。

Answer 1

我找到了一個解決方案，它實際上直接來自 joblib（雖然很難找到）。

以下將不僅使用大約 20% 的單個 CPU 內核執行。 最重要的是，它更快（110ms vs 140ms 在我的機器上進行一次迭代，睡眠時間為 100ms。計算時間減少了 4 倍）。 我猜這是因為要處理的音頻部分相對較小，並且並行化的開銷更大。

from threadpoolctl import threadpool_limits
while True:
    with threadpool_limits(limits=1):
        librosa.feature.mfcc(y=np.random.normal(0, 1, (26783, )), sr=16000, n_mfcc=24, win_length=640, hop_length=640, n_fft=640)
    time.sleep(0.1)