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 distutils import version 22 23from tensorflow.compiler.tf2tensorrt import _pywrap_py_utils 24from tensorflow.core.protobuf import rewriter_config_pb2 25 26 27def disable_non_trt_optimizers_in_rewriter_config(rewriter_config): 28 """Modifies rewriter_config to disable all non-TRT optimizations.""" 29 off = rewriter_config_pb2.RewriterConfig.OFF 30 31 rewriter_config.arithmetic_optimization = off 32 rewriter_config.auto_mixed_precision = off 33 rewriter_config.auto_parallel.enable = False 34 rewriter_config.constant_folding = off 35 rewriter_config.debug_stripper = off 36 rewriter_config.dependency_optimization = off 37 # This one needs to be ON to allow TF-TRT 38 rewriter_config.disable_meta_optimizer = False 39 rewriter_config.disable_model_pruning = True 40 rewriter_config.function_optimization = off 41 rewriter_config.implementation_selector = off 42 rewriter_config.layout_optimizer = off 43 rewriter_config.loop_optimization = off 44 rewriter_config.memory_optimization = ( 45 rewriter_config_pb2.RewriterConfig.NO_MEM_OPT) 46 rewriter_config.min_graph_nodes = -1 47 rewriter_config.pin_to_host_optimization = off 48 rewriter_config.remapping = off 49 rewriter_config.scoped_allocator_optimization = off 50 rewriter_config.shape_optimization = off 51 52 53def _version_tuple_to_string(ver_tuple): 54 assert isinstance(ver_tuple, tuple) 55 assert len(ver_tuple) == 3 56 57 ver_tuple = [str(x) for x in ver_tuple] 58 return ".".join(ver_tuple) 59 60 61def _is_tensorrt_version_greater_equal(trt_ver, target_ver): 62 trt_ver = version.LooseVersion(_version_tuple_to_string(trt_ver)) 63 target_ver = version.LooseVersion(_version_tuple_to_string(target_ver)) 64 65 return trt_ver >= target_ver 66 67 68def is_linked_tensorrt_version_greater_equal(major, minor=0, patch=0): 69 ver = _pywrap_py_utils.get_linked_tensorrt_version() 70 return _is_tensorrt_version_greater_equal(ver, (major, minor, patch)) 71 72 73def is_loaded_tensorrt_version_greater_equal(major, minor=0, patch=0): 74 ver = _pywrap_py_utils.get_loaded_tensorrt_version() 75 return _is_tensorrt_version_greater_equal(ver, (major, minor, patch)) 76