一、限制pytorch 运行的线程数

假如我有4个cpu ,但是只想让Pytorch在1个cpu上运行

import os

cpu_num = 1 # 这里设置成你想运行的CPU个数
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)

假如我有80个cpu ,一般pytorch 默认是使用一半的CPU运行的(40个核),那想让pytorch在所有CPU上运行!!!(小声bb, 没有GPU,我还不能满核吗!),该怎么办呢?

直接使用下边代码

import os
from multiprocessing import cpu_count

cpu_num = cpu_count() # 自动获取最大核心数目
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)

看!!! 80个核了!

在这里插入图片描述

----来自没有GPU的穷孩子。。。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐