一、限制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的穷孩子。。。