Adds support for SavedModel to TOCO.

PiperOrigin-RevId: 190242571

1 files changed, 5 insertions, 2 deletions

diff --git a/tensorflow/contrib/lite/toco/args.h b/tensorflow/contrib/lite/toco/args.h
index 59a6115920..7b71792ff7 100644
--- a/tensorflow/contrib/lite/toco/args.h
+++ b/tensorflow/contrib/lite/toco/args.h
@@ -190,6 +190,7 @@ struct ParsedModelFlags {
   Arg<string> output_array;
   Arg<string> output_arrays;
   Arg<string> input_shapes;
+  Arg<int> batch_size = Arg<int>(1);
   Arg<float> mean_value = Arg<float>(0.f);
   Arg<string> mean_values;
   Arg<float> std_value = Arg<float>(1.f);
@@ -215,9 +216,11 @@ struct ParsedModelFlags {
 // you want). See toco_cmdline_flags.cc for details.
 struct ParsedTocoFlags {
   Arg<string> input_file;
+  Arg<string> savedmodel_directory;
   Arg<string> output_file;
-  Arg<string> input_format;
-  Arg<string> output_format;
+  Arg<string> input_format = Arg<string>("TENSORFLOW_GRAPHDEF");
+  Arg<string> output_format = Arg<string>("TFLITE");
+  Arg<string> savedmodel_tagset;
   // TODO(aselle): command_line_flags  doesn't support doubles
   Arg<float> default_ranges_min = Arg<float>(0.);
   Arg<float> default_ranges_max = Arg<float>(0.);