1# Copyright 2020 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"""Exposes the Python wrapper conversion to trt_graph.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21from tensorflow.core.protobuf import rewriter_config_pb2 22 23 24def disable_non_trt_optimizers_in_rewriter_config(rewriter_config): 25 """Modifies rewriter_config to disable all non-TRT optimizations.""" 26 off = rewriter_config_pb2.RewriterConfig.OFF 27 28 rewriter_config.arithmetic_optimization = off 29 rewriter_config.auto_mixed_precision = off 30 rewriter_config.auto_parallel.enable = False 31 rewriter_config.constant_folding = off 32 rewriter_config.debug_stripper = off 33 rewriter_config.dependency_optimization = off 34 # This one needs to be ON to allow TF-TRT 35 rewriter_config.disable_meta_optimizer = False 36 rewriter_config.disable_model_pruning = True 37 rewriter_config.function_optimization = off 38 rewriter_config.implementation_selector = off 39 rewriter_config.layout_optimizer = off 40 rewriter_config.loop_optimization = off 41 rewriter_config.memory_optimization = ( 42 rewriter_config_pb2.RewriterConfig.NO_MEM_OPT) 43 rewriter_config.min_graph_nodes = -1 44 rewriter_config.pin_to_host_optimization = off 45 rewriter_config.remapping = off 46 rewriter_config.scoped_allocator_optimization = off 47 rewriter_config.shape_optimization = off 48