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# ============================================================================== 15r"""Prints a header file to be used with SELECTIVE_REGISTRATION. 16 17An 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 22Then when compiling tensorflow, include ops_to_register.h in the include search 23path and pass -DSELECTIVE_REGISTRATION and -DSUPPORT_SELECTIVE_REGISTRATION 24 - see core/framework/selective_registration.h for more details. 25 26When 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 34from __future__ import absolute_import 35from __future__ import division 36from __future__ import print_function 37 38import argparse 39import sys 40 41from tensorflow.python.platform import app 42from tensorflow.python.tools import selective_registration_header_lib 43 44FLAGS = None 45 46 47def 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 53if __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