1# Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Python console command to invoke TOCO from serialized protos.""" 16import argparse 17import sys 18 19# We need to import pywrap_tensorflow prior to the toco wrapper. 20# pylint: disable=invalid-import-order,g-bad-import-order 21from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import 22from tensorflow.python import _pywrap_toco_api 23from absl import app 24 25FLAGS = None 26 27 28def execute(unused_args): 29 """Runs the converter.""" 30 with open(FLAGS.model_proto_file, "rb") as model_file: 31 model_str = model_file.read() 32 33 with open(FLAGS.toco_proto_file, "rb") as toco_file: 34 toco_str = toco_file.read() 35 36 with open(FLAGS.model_input_file, "rb") as input_file: 37 input_str = input_file.read() 38 39 debug_info_str = None 40 if FLAGS.debug_proto_file: 41 with open(FLAGS.debug_proto_file, "rb") as debug_info_file: 42 debug_info_str = debug_info_file.read() 43 44 enable_mlir_converter = FLAGS.enable_mlir_converter 45 46 output_str = _pywrap_toco_api.TocoConvert( 47 model_str, 48 toco_str, 49 input_str, 50 False, # extended_return 51 debug_info_str, 52 enable_mlir_converter) 53 open(FLAGS.model_output_file, "wb").write(output_str) 54 sys.exit(0) 55 56 57def main(): 58 global FLAGS 59 parser = argparse.ArgumentParser( 60 description="Invoke toco using protos as input.") 61 parser.add_argument( 62 "model_proto_file", 63 type=str, 64 help="File containing serialized proto that describes the model.") 65 parser.add_argument( 66 "toco_proto_file", 67 type=str, 68 help="File containing serialized proto describing how TOCO should run.") 69 parser.add_argument( 70 "model_input_file", type=str, help="Input model is read from this file.") 71 parser.add_argument( 72 "model_output_file", 73 type=str, 74 help="Result of applying TOCO conversion is written here.") 75 parser.add_argument( 76 "--debug_proto_file", 77 type=str, 78 default="", 79 help=("File containing serialized `GraphDebugInfo` proto that describes " 80 "logging information.")) 81 parser.add_argument( 82 "--enable_mlir_converter", 83 action="store_true", 84 help=("Boolean indicating whether to enable MLIR-based conversion " 85 "instead of TOCO conversion. (default False)")) 86 87 FLAGS, unparsed = parser.parse_known_args() 88 89 app.run(main=execute, argv=[sys.argv[0]] + unparsed) 90 91 92if __name__ == "__main__": 93 main() 94