Factor parsing the args in optimize_for_inference to a separate function, to allow reusing the binary.

Change: 147666737

1 files changed, 7 insertions, 2 deletions

diff --git a/tensorflow/python/tools/optimize_for_inference.py b/tensorflow/python/tools/optimize_for_inference.py
index 4fa1ac5042..902748d55e 100644
--- a/tensorflow/python/tools/optimize_for_inference.py
+++ b/tensorflow/python/tools/optimize_for_inference.py
@@ -99,7 +99,8 @@ def main(unused_args):
   return 0
 
 
-if __name__ == "__main__":
+def parse_args():
+  """Parses command line arguments."""
   parser = argparse.ArgumentParser()
   parser.register("type", "bool", lambda v: v.lower() == "true")
   parser.add_argument(
@@ -137,5 +138,9 @@ if __name__ == "__main__":
       type=int,
       default=dtypes.float32.as_datatype_enum,
       help="The AttrValue enum to use for placeholders.")
-  FLAGS, unparsed = parser.parse_known_args()
+  return parser.parse_known_args()
+
+
+if __name__ == "__main__":
+  FLAGS, unparsed = parse_args()
   app.run(main=main, argv=[sys.argv[0]] + unparsed)