Merge pull request #725 from Intel-tensorflow/nhasabni/t2t-cpu-optimizations

Enabling setting Tensorflow inter_op and intra_op

Author: lukaszkaiser

tensor2tensor/bin/t2t_trainer.py

@@ -57,6 +57,10 @@
 flags.DEFINE_string("tmp_dir", "/tmp/t2t_datagen",
                     "Temporary storage directory, used if --generate_data.")
 flags.DEFINE_bool("profile", False, "Profile performance?")
+flags.DEFINE_integer("inter_op_parallelism_threads", 0, "Number of inter_op_parallelism_threads "
+                     "to use for CPU. See TensorFlow config.proto for details.")
+flags.DEFINE_integer("intra_op_parallelism_threads", 0, "Number of intra_op_parallelism_threads "
+                     "to use for CPU. See TensorFlow config.proto for details.")
 
 # To maintain compatibility with some internal libs, we guard against these flag
 # definitions possibly erring. Apologies for the ugliness.
@@ -206,7 +210,9 @@ def create_run_config(hp):
       worker_id=FLAGS.worker_id,
       worker_job=FLAGS.worker_job,
       random_seed=FLAGS.random_seed,
-      tpu_infeed_sleep_secs=FLAGS.tpu_infeed_sleep_secs)
+      tpu_infeed_sleep_secs=FLAGS.tpu_infeed_sleep_secs,
+      inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
+      intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads)
 
 
 def generate_data():

tensor2tensor/utils/trainer_lib.py

@@ -40,7 +40,9 @@
 def create_session_config(log_device_placement=False,
                           enable_graph_rewriter=False,
                           gpu_mem_fraction=0.95,
-                          use_tpu=False):
+                          use_tpu=False,
+                          inter_op_parallelism_threads=0,
+                          intra_op_parallelism_threads=0):
   """The TensorFlow Session config to use."""
   if use_tpu:
     graph_options = tf.GraphOptions()
@@ -60,7 +62,9 @@ def create_session_config(log_device_placement=False,
       allow_soft_placement=True,
       graph_options=graph_options,
       gpu_options=gpu_options,
-      log_device_placement=log_device_placement)
+      log_device_placement=log_device_placement,
+      inter_op_parallelism_threads=inter_op_parallelism_threads,
+      intra_op_parallelism_threads=intra_op_parallelism_threads)
   return config
 
 
@@ -108,13 +112,17 @@ def create_run_config(master="",
                       random_seed=None,
                       sync=False,
                       tpu_infeed_sleep_secs=None,
-                      use_tpu=False):
+                      use_tpu=False,
+                      inter_op_parallelism_threads=0,
+                      intra_op_parallelism_threads=0):
   """Create RunConfig, TPUConfig, and Parallelism object."""
   session_config = create_session_config(
       log_device_placement=log_device_placement,
       enable_graph_rewriter=enable_graph_rewriter,
       gpu_mem_fraction=gpu_mem_fraction,
-      use_tpu=use_tpu)
+      use_tpu=use_tpu,
+      inter_op_parallelism_threads=inter_op_parallelism_threads,
+      intra_op_parallelism_threads=intra_op_parallelism_threads)
   run_config_args = {
       "master": master,
       "evaluation_master": master,