如何解决PyTorch程序占用较高CPU问题

PyTorch程序占用较高CPU问题

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

解决办法

在main程序中加以下代码：

torch.set_num_threads(1)

设置最多使用1个CPU核。

问题解决~

pytorch训练占用cpu过高，num_works和set_num_threads设置均无效，发现是数据扩增问题

这是我使用pytorch训练模型的时候，出现cpu占用过多的情况，无关pytorch版本

dataloader 的 num_work=1 的时候单线程cpu占用量2800，也就是一半的cpu，我服务器一共28*2个逻辑cpu

dataloader 的 num_work=8 的时候8个线程cpu占用500-700，合计2800

使用网上教程的如下指令，没有任何效果

    cpu_num = 1
    os.environ["OMP_NUM_THREADS"] = str(cpu_num)
    os.environ["OPENBLAS_NUM_THREADS"] = str(cpu_num)
    os.environ["MKL_NUM_THREADS"] = str(cpu_num)
    os.environ["VECLIB_MAXIMUM_THREADS"] = str(cpu_num)
    os.environ["NUMEXPR_NUM_THREADS"] = str(cpu_num)
    torch.set_num_threads(cpu_num)

最后发现是数据扩增惹的祸，我是用的数据扩增库是albumentations，其中所使用的基于opencv的各种数据变换占用了太多的cpu，因此只需要限制opencv的线程数目则可：

import cv2
import albumentations
cv2.setNumThreads(1)	# 0也可以

再设置 dataloader 的 num_work=2 的情况

虽然不清楚这里为什么有3个，但目的达成了，而且训练速度和cpu占用2800的时候一样。

在albumentations的github网址也有相关联的一小段说明，但不是本篇文章遇到的问题， cv2.ocl.setUseOpenCL(False) 对于本篇文章问题可以不加。

总结

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