如何在PyTorch中禁用多线程？

Question

我正在尝试确保用C++构建的PyTorch程序只使用单个线程。该程序在CPU上运行。

它具有相当小的模型，多线程并没有帮助，实际上会导致问题，因为我的程序已经是多线程的。我已经调用了：

    at::set_num_interop_threads(1);
    at::set_num_threads(1);
    torch::set_num_threads(1);
    omp_set_num_threads(1);
    omp_set_dynamic(0);
    omp_set_nested(0);

此外，我已将环境变量OPENBLAS_NUM_THREADS设置为1。

然而，即使我在单线程中进行孵化，htop上仍会显示16个线程，并且机器的16个处理器都会达到100％。

我是否遗漏了什么？什么原因呢？

- willem

Answer 1

从PyTorch文档中，可以执行以下操作：

torch.set_num_threads(1)

为了保险起见，在实例化任何模型之前（即在导入后立即执行此操作）。这对我有用。

更多信息：https://jdhao.github.io/2020/07/06/pytorch_set_num_threads/

- mxbi