PyTorch如何限制CPU线程数目

PyTorch限制CPU线程数目

终端命令 htop ：查看 CPU 利用率

限制 python 的cpu 利用率：

import os
cpu_num = 4 # 这里设置成你想运行的CPU个数
os.environ["OMP_NUM_THREADS"] = str(cpu_num)  # noqa
os.environ["MKL_NUM_THREADS"] = str(cpu_num) # noqa
torch.set_num_threads(cpu_num )

原因

tensor 没有 .to('cuda') ,tensor 还是部署在 cpu 上，耗的是 cpu 的资源，导致 cpu 占用率过高。

在读数据的时候设置 num_workers = 0，也可能会导致 cpu 占用率过高。

PyTorch程序占用较高CPU解决

在使用服务器运行PyTorch程序时，使用TOP命令查看CPU状态，发现CPU占用率竟然达到2000%，已经影响到其他人正常使用服务器。

解决办法

在main程序中加以下代码：

torch.set_num_threads(1)

设置最多使用1个CPU核。问题解决~

总结

以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。