diff --git a/darkflow/net/build.py b/darkflow/net/build.py index 1359f9f12..c78149cb6 100644 --- a/darkflow/net/build.py +++ b/darkflow/net/build.py @@ -141,6 +141,9 @@ def setup_meta_ops(self): if self.FLAGS.summary: self.summary_op = tf.summary.merge_all() self.writer = tf.summary.FileWriter(self.FLAGS.summary + 'train') + if self.FLAGS.cpus: + cfg['inter_op_parallelism_threads'] = self.FLAGS.cpus + cfg['intra_op_parallelism_threads'] = self.FLAGS.cpus self.sess = tf.Session(config = tf.ConfigProto(**cfg)) self.sess.run(tf.global_variables_initializer()) @@ -174,4 +177,4 @@ def savepb(self): json.dump(self.meta, fp) self.say('Saving const graph def to {}'.format(name)) graph_def = tfnet_pb.sess.graph_def - tf.train.write_graph(graph_def,'./', name, False) \ No newline at end of file + tf.train.write_graph(graph_def,'./', name, False)