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

Question

根据this document，我必须插入os.environ["MKL_NUM_THREADS"] = "1"以防止numpy使用多个内核。但是我应该在哪里插入这个？在 numpy 之前或在每个 actor 的构造中？应该用在 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 的快照，我在 main.py 的开头插入 os.environ["MKL_NUM_THREADS"] = "1"。它有效吗？

Answer 1

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

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

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

启动 Python 脚本来实现相同的目的

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

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

where should I insert `os.environ["MKL_NUM_THREADS"] = "1"`?

python

numpy

ray