Keras / tensorflow - limit number of cores (intra_op_parallelism_threads not working)

Question

I've been trying to run keras on a CPU cluster, and for this I need to limit the number of cores used (it's a shared system). So to limit the number of cores, I landed on this answer. However, this simply doesn't work. I tried running with this basic code:

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)

When I run this and check htop, it uses all (128) logical cores. Is this a bug in keras? Or am I doing something wrong?

Keras says that my CPU supports SSE4.1 and SSE4.2, which are not used because I didn't compile from binary. Will compiling from binary also fix the original question?

EDIT: I've found a workaround when launching the keras script from a unix machine:

taskset -c 0-23 python keras_script.py

This will run the script on the first 24 cores of the machine. It works, but it would still be nice if this was available from within keras/tensorflow.

Answer 1

I found this snippet of code that works for me, hope it helps:

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)