Add Inter/Intra_op_parallelism_threads Support to Wide and Deep (#5046)

official/wide_deep/README.md

@@ -39,6 +39,14 @@ To run the *wide* or *deep*-only models, set the `--model_type` flag to `wide` o
 
 The final accuracy should be over 83% with any of the three model types.
 
+You can also experiment with `-inter` and `-intra` flag to explore inter/intra op parallelism for potential better performance as follows:
+
+```
+python census_main.py --inter=<int> --intra=<int>
+```
+Please note the above optional inter/intra op does not affect model accuracy. These are TensorFlow framework configurations that only affect execution time.
+For more details regarding the above inter/intra flags, please refer to [Optimizing_for_CPU](https://www.tensorflow.org/performance/performance_guide#optimizing_for_cpu) or [TensorFlow config.proto source code](https://github.com/tensorflow/tensorflow/blob/26b4dfa65d360f2793ad75083c797d57f8661b93/tensorflow/core/protobuf/config.proto#L165).
+
 ### TensorBoard
 
 Run TensorBoard to inspect the details about the graph and training progression.

official/wide_deep/census_main.py

@@ -33,10 +33,12 @@ def define_census_flags():
                           model_dir='/tmp/census_model',
                           train_epochs=40,
                           epochs_between_evals=2,
+                          inter_op_parallelism_threads=0,
+                          intra_op_parallelism_threads=0,
                           batch_size=40)
 
 
-def build_estimator(model_dir, model_type, model_column_fn):
+def build_estimator(model_dir, model_type, model_column_fn, inter_op, intra_op):
   """Build an estimator appropriate for the given model type."""
   wide_columns, deep_columns = model_column_fn()
   hidden_units = [100, 75, 50, 25]
@@ -44,7 +46,9 @@ def build_estimator(model_dir, model_type, model_column_fn):
   # Create a tf.estimator.RunConfig to ensure the model is run on CPU, which
   # trains faster than GPU for this model.
   run_config = tf.estimator.RunConfig().replace(
-      session_config=tf.ConfigProto(device_count={'GPU': 0}))
+      session_config=tf.ConfigProto(device_count={'GPU': 0},
+                                    inter_op_parallelism_threads=inter_op,
+                                    intra_op_parallelism_threads=intra_op))
 
   if model_type == 'wide':
     return tf.estimator.LinearClassifier(

official/wide_deep/movielens_main.py

@@ -44,6 +44,8 @@ def define_movie_flags():
                           model_type="deep",
                           train_epochs=50,
                           epochs_between_evals=5,
+                          inter_op_parallelism_threads=0,
+                          intra_op_parallelism_threads=0,
                           batch_size=256)
 
   @flags.validator("stop_threshold",
@@ -52,13 +54,17 @@ def define_movie_flags():
     return stop_threshold is None
 
 
-def build_estimator(model_dir, model_type, model_column_fn):
+def build_estimator(model_dir, model_type, model_column_fn, inter_op, intra_op):
   """Build an estimator appropriate for the given model type."""
   if model_type != "deep":
     raise NotImplementedError("movie dataset only supports `deep` model_type")
   _, deep_columns = model_column_fn()
   hidden_units = [256, 256, 256, 128]
 
+  run_config = tf.estimator.RunConfig().replace(
+      session_config=tf.ConfigProto(device_count={'GPU': 0},
+                                    inter_op_parallelism_threads=inter_op,
+                                    intra_op_parallelism_threads=intra_op))
   return tf.estimator.DNNRegressor(
       model_dir=model_dir,
       feature_columns=deep_columns,

official/wide_deep/wide_deep_run_loop.py

@@ -48,6 +48,14 @@ def define_wide_deep_flags():
   flags.DEFINE_boolean(
       name="download_if_missing", default=True, help=flags_core.help_wrap(
           "Download data to data_dir if it is not already present."))
+  flags.DEFINE_integer(
+      name="inter_op_parallelism_threads", short_name="inter", default=0,
+      help="Number of threads to use for inter-op parallelism. "
+           "If left as default value of 0, the system will pick an appropriate number.")
+  flags.DEFINE_integer(
+      name="intra_op_parallelism_threads", short_name="intra", default=0,
+      help="Number of threads to use for intra-op parallelism. "
+           "If left as default value of 0, the system will pick an appropriate number.")
 
 
 def export_model(model, model_type, export_dir, model_column_fn):
@@ -78,7 +86,9 @@ def run_loop(name, train_input_fn, eval_input_fn, model_column_fn,
   model_helpers.apply_clean(flags.FLAGS)
   model = build_estimator_fn(
       model_dir=flags_obj.model_dir, model_type=flags_obj.model_type,
-      model_column_fn=model_column_fn)
+      model_column_fn=model_column_fn,
+      inter_op=flags_obj.inter_op_parallelism_threads,
+      intra_op=flags_obj.intra_op_parallelism_threads)
 
   run_params = {
       'batch_size': flags_obj.batch_size,