1 # Copyright 2016 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 r"""Prints a header file to be used with SELECTIVE_REGISTRATION. 16 17 An example of command-line usage is: 18 bazel build tensorflow/python/tools:print_selective_registration_header && \ 19 bazel-bin/tensorflow/python/tools/print_selective_registration_header \ 20 --graphs=path/to/graph.pb > ops_to_register.h 21 22 Then when compiling tensorflow, include ops_to_register.h in the include search 23 path and pass -DSELECTIVE_REGISTRATION and -DSUPPORT_SELECTIVE_REGISTRATION 24 - see core/framework/selective_registration.h for more details. 25 26 When compiling for Android: 27 bazel build -c opt --copt="-DSELECTIVE_REGISTRATION" \ 28 --copt="-DSUPPORT_SELECTIVE_REGISTRATION" \ 29 //tensorflow/contrib/android:libtensorflow_inference.so \ 30 --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ 31 --crosstool_top=//external:android/crosstool --cpu=armeabi-v7a 32 """ 33 34 from __future__ import absolute_import 35 from __future__ import division 36 from __future__ import print_function 37 38 import argparse 39 import sys 40 41 from tensorflow.python.platform import app 42 from tensorflow.python.tools import selective_registration_header_lib 43 44 FLAGS = None 45 46 47 def main(unused_argv): 48 graphs = FLAGS.graphs.split(',') 49 print(selective_registration_header_lib.get_header( 50 graphs, FLAGS.proto_fileformat, FLAGS.default_ops)) 51 52 53 if __name__ == '__main__': 54 parser = argparse.ArgumentParser() 55 parser.register('type', 'bool', lambda v: v.lower() == 'true') 56 parser.add_argument( 57 '--graphs', 58 type=str, 59 default='', 60 help='Comma-separated list of paths to model files to be analyzed.', 61 required=True) 62 parser.add_argument( 63 '--proto_fileformat', 64 type=str, 65 default='rawproto', 66 help='Format of proto file, either textproto or rawproto.') 67 parser.add_argument( 68 '--default_ops', 69 type=str, 70 default='NoOp:NoOp,_Recv:RecvOp,_Send:SendOp', 71 help='Default operator:kernel pairs to always include implementation for.' 72 'Pass "all" to have all operators and kernels included; note that this ' 73 'should be used only when it is useful compared with simply not using ' 74 'selective registration, as it can in some cases limit the effect of ' 75 'compilation caches') 76 77 FLAGS, unparsed = parser.parse_known_args() 78 app.run(main=main, argv=[sys.argv[0]] + unparsed) 79
