tensorflow - 限制使用的核心数量

Question

我试图限制 tf session 使用的核心数量，但它不起作用。这就是我初始化 session 的方式:

sess = tf.Session(config=tf.ConfigProto(inter_op_parallelism_threads=1,
                                        intra_op_parallelism_threads=1,
                                        use_per_session_threads=True))

系统有 12 个核心/24 个线程，我可以看到其中 40-60% 在任何给定时间点都在使用。该系统也有 8 个 GPU，但我使用 tf.device('/cpu:0') 构建整个图表。

更新:澄清一下，该图本身是一个简单的 LSTM-RNN，它与 tf 源代码中的示例非常接近。为了完整起见，这里是完整的图表:

node_input  = tf.placeholder(tf.float32, [n_steps, batch_size, input_size],  name = 'input')
list_input  = [tf.reshape(i, (batch_size, input_size)) for i in tf.split(0, n_steps, node_input)]
node_target = tf.placeholder(tf.float32, [n_steps, batch_size, output_size], name = 'target')
node_target_flattened = tf.reshape(tf.transpose(node_target, perm = [1, 0, 2]), [-1, output_size])
node_max_length = tf.placeholder(tf.int32, name = 'batch_max_length')
node_cell_initializer = tf.random_uniform_initializer(-0.1, 0.1) 
node_cell = LSTMCell(state_size, input_size, initializer = node_cell_initializer)  
node_initial_state = node_cell.zero_state(batch_size, tf.float32)
nodes_output, nodes_state = rnn(node_cell, 
                                list_input, 
                                initial_state = node_initial_state, 
                                sequence_length = node_max_length)
node_output_flattened = tf.reshape(tf.concat(1, nodes_output), [-1, state_size])
node_softmax_w = tf.Variable(tf.random_uniform([state_size, output_size]), name = 'softmax_w')
node_softmax_b = tf.Variable(tf.zeros([output_size]), name = 'softmax_b')
node_logit = tf.matmul(node_output_flattened, node_softmax_w) + node_softmax_b
node_cross_entropy = tf.nn.softmax_cross_entropy_with_logits(node_logit, node_target_flattened, name = 'cross_entropy')
node_loss = tf.reduce_mean(node_cross_entropy, name = 'loss')
node_optimizer = tf.train.AdamOptimizer().minimize(node_loss)
node_op_initializer = tf.initialize_all_variables()

需要注意的一件重要事情是，如果我第一次调用 tf.Session，我传入了适当的参数，那么 session 确实仅在单个核。问题是，在后续运行中，我无法更改行为，即使我使用了 use_per_session_threads ，它应该专门允许特定于 session 的设置。 IE。即使我使用 sess.close() 关闭 session 并使用新选项启动一个新 session ，原始行为仍然保持不变，除非我重新启动 python 内核(这是非常昂贵的，因为它需要花费近小时加载我的数据)。

Answer 1

use_per_session_threads 只会影响 inter_op_parallelism_threads，但不会影响 intra_op_parallelism_threads。 intra_op_parallelism_threads 将用于 Eigen 线程池(请参阅 here )，它始终是全局的，因此后续 session 将不再影响它。

请注意，还有其他 TF 函数也可以触发 Eigen 线程池的初始化，因此在创建第一个 tf.Session 之前它可能已经初始化。一个示例是 tensorflow.python.client.device_lib.list_local_devices()。

我在 Python 脚本的早期就用一种方法解决了这个问题，我使用适当的值创建了一个虚拟 session 。