代码之家  ›  专栏  ›  技术社区  ›  wouterdobbels

Keras/tensorflow-限制芯数(intra\u op\u parallelism\u threads not working)

  •  8
  • wouterdobbels  · 技术社区  · 7 年前

    我一直在尝试在CPU集群上运行keras,为此我需要限制使用的内核数量(这是一个共享系统)。为了限制内核的数量,我在 this answer .然而,这根本不起作用。我尝试使用以下基本代码运行:

    from keras.applications.vgg16 import VGG16
    from keras import backend as K
    import numpy as np
    
    conf = K.tf.ConfigProto(device_count={'CPU': 1}, 
                            intra_op_parallelism_threads=2, 
                            inter_op_parallelism_threads=2)
    K.set_session(K.tf.Session(config=conf))
    model = VGG16(weights='imagenet', include_top=False)
    x = np.random.randn(1000, 224, 224, 3)
    features = model.predict(x)
    

    当我运行这个并检查 htop ,它使用所有(128)个逻辑核。这是keras中的虫子吗?还是我做错了什么?

    Keras说我的CPU支持SSE4。1和SSE4。2,因为我没有从二进制编译,所以没有使用。从二进制文件编译也会解决原来的问题吗?

    编辑 :我在从unix计算机启动keras脚本时找到了一个解决方法:

    taskset -c 0-23 python keras_script.py
    

    这将在机器的前24个内核上运行脚本。这是可行的,但如果可以从keras/tensorflow内部获得这一功能,那就更好了。

    1 回复  |  直到 7 年前
        1
  •  1
  •   Pau    6 年前

    我找到了这段适合我的代码,希望它能帮助我:

    from keras import backend as K
    import tensorflow as tf
    jobs = 2 # it means number of cores
    config = tf.ConfigProto(intra_op_parallelism_threads=jobs,
                             inter_op_parallelism_threads=jobs,
                             allow_soft_placement=True,
                             device_count={'CPU': jobs})
    session = tf.Session(config=config)
    K.set_session(session)