python - 我应该在哪里插入`os.environ["MKL_NUM_THREADS"] = "1"`？

Question

根据这个文件，我必须插入os.environ["MKL_NUM_THREADS"] = "1"以防止numpy使用多个核心。但是我应该在哪里插入这个？numpy在每个演员的构建之前还是之中？应该用在Learner?

例如，如果我有以下三个文件，我应该在哪里插入？

""" learner.py """
import numpy as np
import ray

@ray.remote
class Learner:
    def __init__(self):
        # do something

""" worker.py """
import numpy as np
import ray

@ray.remote
class Worker:
    def __init__(self):
        # do something

""" main.py """
# program starts from here

import numpy as np
import ray

from learner import Learner
from worker import Worker

ray.init()

learner = Learner.remote()
worker = Worker.remote()

# do something...

如何检查它是否按预期工作？top这是我真实项目的快照，我os.environ["MKL_NUM_THREADS"] = "1"在main.py. 它有效吗？

Answer 1

您可以通过查看top并确保没有一个工作进程 CPU 利用率超过 100% 来检查它是否正常工作。

您想在调用之前设置它，ray.init()以便在工作进程被派生时，环境变量会传播给它们。

您可以通过使用MKL_NUM_THREADS=1 python script.py.

请注意，此环境变量在使用 MKL 时有意义numpy，但根据 BLAS 实现，您可能需要设置一个环境变量，如OPENBLAS_NUM_THREADS=1or OMP_NUM_THREADS=1。