Skip to content

GitLab

Commit 7905fc4f authored by Jin Young Sohn's avatar Jin Young Sohn Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 362271822
parent 82b922b4
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 0 deletions
official/vision/image_classification/mnist_main.py
View file @ 7905fc4f
...@@ -81,6 +81,9 @@ def run(flags_obj, datasets_override=None, strategy_override=None): ...@@ -81,6 +81,9 @@ def run(flags_obj, datasets_override=None, strategy_override=None):
Returns: Returns:
Dictionary of training and eval stats. Dictionary of training and eval stats.
""" """
# Start TF profiler server.
tf.profiler.experimental.server.start(flags_obj.profiler_port)
strategy = strategy_override or distribute_utils.get_distribution_strategy( strategy = strategy_override or distribute_utils.get_distribution_strategy(
distribution_strategy=flags_obj.distribution_strategy, distribution_strategy=flags_obj.distribution_strategy,
num_gpus=flags_obj.num_gpus, num_gpus=flags_obj.num_gpus,
...@@ -156,6 +159,8 @@ def define_mnist_flags(): ...@@ -156,6 +159,8 @@ def define_mnist_flags():
flags_core.define_distribution() flags_core.define_distribution()
flags.DEFINE_bool('download', False, flags.DEFINE_bool('download', False,
'Whether to download data to `--data_dir`.') 'Whether to download data to `--data_dir`.')
flags.DEFINE_integer('profiler_port', 9012,
'Port to start profiler server on.')
FLAGS.set_default('batch_size', 1024) FLAGS.set_default('batch_size', 1024)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment