Pytorch 占用cpu资源过多

Pytorch 占用cpu资源过多

最近在用pytorch进行一些AI的工作时发现 在部分环境下其cpu占用量过高。

如下：这里只简单的推理了一个小网络且还是用GPU运算的，其cpu占用已经高达1348%了。

top

经过调试发现 主要是将数据从GPU下载到CPU上消耗资源比较高。

解决办法

遇到这个问题，限制一下线程数量就可以了；

torch.set_num_threads(num)

这个方法改变CPU的占用量比较明显。

效果

限制之后的效果如下：CPU使用率明显降下来了。