chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s

This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
@@ -0,0 +1,587 @@
# Description:
# Wrap NVIDIA TensorRT (http://developer.nvidia.com/tensorrt) with tensorflow
# and provide TensorRT operators and converter package.
# APIs are meant to change over time.
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
exports_files(
glob(["*test.py"]),
)
py_library(
name = "test_utils",
srcs = ["test_utils.py"],
strict_deps = True,
)
py_library(
name = "tf_trt_integration_test_base_srcs",
srcs = ["tf_trt_integration_test_base.py"],
strict_deps = True,
visibility = ["//tensorflow/python/compiler/tensorrt:__pkg__"],
deps = [
"//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/compiler/tensorrt:trt_convert_py",
"//tensorflow/python/compiler/tensorrt:utils",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:graph_io",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler:trace",
"//tensorflow/python/saved_model:builder",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:loader",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/saved_model:signature_constants",
"//tensorflow/python/saved_model:signature_def_utils",
"//tensorflow/python/saved_model:tag_constants",
"//tensorflow/python/saved_model:utils",
"//tensorflow/python/tools:saved_model_utils",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
filegroup(
name = "trt_convert_test_data",
srcs = [
"testdata/tf_readvariableop_saved_model/saved_model.pb",
"testdata/tf_readvariableop_saved_model/variables/variables.data-00000-of-00001",
"testdata/tf_readvariableop_saved_model/variables/variables.index",
"testdata/tf_variablev2_saved_model/saved_model.pb",
"testdata/tf_variablev2_saved_model/variables/variables.data-00000-of-00001",
"testdata/tf_variablev2_saved_model/variables/variables.index",
"testdata/tftrt_2.0_saved_model/saved_model.pb",
"testdata/tftrt_2.0_saved_model/variables/variables.data-00000-of-00002",
"testdata/tftrt_2.0_saved_model/variables/variables.data-00001-of-00002",
"testdata/tftrt_2.0_saved_model/variables/variables.index",
],
visibility = ["//tensorflow/python/compiler/tensorrt:__pkg__"],
)
base_tags = [
"no_cuda_on_cpu_tap",
"no_rocm",
"no_windows",
"nomac",
# TODO(b/303453873): Re-enable tests once TensorRT has been updated
"notap",
]
cuda_py_strict_test(
name = "base_test",
srcs = ["base_test.py"],
tags = base_tags + [
"no_oss", # TODO(b/165611343): Need to address the failures for CUDA 11 in OSS build.
"notap", # TODO(b/291653036): Some of the base tests are failing w/ TensorRT-8
],
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/compiler/tensorrt:utils",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "conv2d_test",
srcs = ["conv2d_test.py"],
tags = base_tags + ["no_oss"], # b/198501457
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "combined_nms_test",
srcs = ["combined_nms_test.py"],
tags = base_tags + ["no_oss"],
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/compiler/tensorrt:utils",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:image_ops_impl",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "tf_function_test",
srcs = ["tf_function_test.py"],
tags = base_tags + [
"manual", # TODO(b/231239602): re-enable once naming issue is resolved.
"no_oss", # TODO(b/231239602): re-enable once naming issue is resolved.
"notap", # TODO(b/231239602): re-enable once naming issue is resolved.
],
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/compiler/tensorrt:trt_convert_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model:constants",
"//tensorflow/python/saved_model:loader",
"//tensorflow/python/util:compat",
],
)
cuda_py_strict_test(
name = "batch_matmul_test",
srcs = ["batch_matmul_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "biasadd_matmul_test",
srcs = ["biasadd_matmul_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "binary_tensor_weight_broadcast_test",
srcs = ["binary_tensor_weight_broadcast_test.py"],
# TODO(b/297490791): Reenable after TensorRT regression has been fixed
tags = base_tags + ["no_oss"],
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:tf_logging",
],
)
cuda_py_strict_test(
name = "bool_test",
srcs = ["bool_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/compiler/tensorrt:utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "cast_test",
srcs = ["cast_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "concatenation_test",
srcs = ["concatenation_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "const_broadcast_test",
srcs = ["const_broadcast_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "data_dependent_shape_test",
srcs = ["data_dependent_shape_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "dynamic_input_shapes_test",
srcs = ["dynamic_input_shapes_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "identity_output_test",
srcs = ["identity_output_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "int32_test",
srcs = ["int32_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "lru_cache_test",
srcs = ["lru_cache_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "memory_alignment_test",
srcs = ["memory_alignment_test.py"],
tags = base_tags + ["no_oss"], # TODO(b/285588022)
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "multi_connection_neighbor_engine_test",
srcs = ["multi_connection_neighbor_engine_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "neighboring_engine_test",
srcs = ["neighboring_engine_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "quantization_test",
srcs = ["quantization_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "rank_two_test",
srcs = ["rank_two_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "reshape_transpose_test",
srcs = ["reshape_transpose_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "topk_test",
srcs = ["topk_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "trt_engine_op_shape_test",
srcs = ["trt_engine_op_shape_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:signature_constants",
"//tensorflow/python/saved_model:tag_constants",
],
)
cuda_py_strict_test(
name = "trt_mode_test",
srcs = ["trt_mode_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "unary_test",
srcs = ["unary_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "vgg_block_nchw_test",
srcs = ["vgg_block_nchw_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_impl",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "vgg_block_test",
srcs = ["vgg_block_test.py"],
tags = base_tags,
xla_enable_strict_auto_jit = False,
deps = [
":tf_trt_integration_test_base_srcs",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_impl",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/platform:client_testlib",
],
)
test_suite(
name = "tf_trt_integration_test",
tests = [
"batch_matmul_test",
"biasadd_matmul_test",
"binary_tensor_weight_broadcast_test",
"bool_test",
"cast_test",
"combined_nms_test",
"concatenation_test",
"const_broadcast_test",
"data_dependent_shape_test",
"dynamic_input_shapes_test",
"identity_output_test",
"int32_test",
"lru_cache_test",
"memory_alignment_test",
"multi_connection_neighbor_engine_test",
"neighboring_engine_test",
"quantization_test",
"rank_two_test",
"reshape_transpose_test",
"tf_function_test",
"topk_test",
"trt_engine_op_shape_test",
"trt_mode_test",
"unary_test",
"vgg_block_nchw_test",
"vgg_block_test",
],
)
test_suite(
name = "tf_trt_integration_test_no_oss",
tests = [
"base_test",
"conv2d_test",
],
)
@@ -0,0 +1,135 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Testing the impact of graph node _tftrt_op_max_batch_size annotation on TRTEngineOp attributes."""
import unittest
from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
class MaxBatchSizesTestBase(trt_test.TfTrtIntegrationTestBase):
@classmethod
def setUpClass(cls):
if cls is MaxBatchSizesTestBase:
raise unittest.SkipTest(
'MaxBatchSizesTestBase defines base class for other tests.')
super(MaxBatchSizesTestBase, cls).setUpClass()
@property
def tensor_shapes(self):
return [[1, 512, 1, 1], [64, 2, 2, 2], [32, 4, 2, 2], [16, 8, 2, 2]]
@property
def max_batch_sizes(self):
return [shape[0] for shape in self.tensor_shapes]
def GetParams(self):
"""Gets the build parameters for the test."""
return self.BuildParams(
self.GraphFn,
dtype=dtypes.float32,
input_shapes=[self.tensor_shapes[0]],
output_shapes=[self.tensor_shapes[-1]])
def ShouldRunTest(self, run_params):
# The maximum batch size for dynamic engines will be the actual batch size
# detected at runtime. Therefore, we don't run the test with dynamic
# engines.
return (not run_params.dynamic_engine, 'test static engine only.')
def GetMaxBatchSize(self, run_params):
"""Returns the max_batch_size that the converter should use for tests."""
if run_params.dynamic_engine:
return None
return min(self.max_batch_sizes)
def ExpectedEnginesToBuild(self, run_params):
"""Checks that the expected engine is built.
Args:
run_params: the run parameters.
Returns:
the expected engines to build.
There shall be engines generated for each maximum batch size.
"""
return [
f'TRTEngineOp_{seq_id:03d}'
for seq_id in range(len(self.max_batch_sizes))
]
def ExpectedMaxBatchSizes(self, run_params):
"""Checks that the expected maximum batch sizes for the generated engines.
Args:
run_params: the run parameters.
Returns:
the expected maximum batch sizes for the generated engines.
There shall be engines generated for each maximum batch size.
"""
return self.max_batch_sizes
class AnnotateMaxBatchSizesTest(MaxBatchSizesTestBase):
def GraphFn(self, inp):
"""Builds a tf.Graph for the test."""
tensor = inp * 2.0
tensor = array_ops.reshape(tensor, [-1] + self.tensor_shapes[1][1:])
with ops.get_default_graph()._attr_scope({
'_tftrt_op_max_batch_size':
attr_value_pb2.AttrValue(i=self.max_batch_sizes[1])
}):
tensor = tensor + 3.0
tensor = array_ops.reshape(tensor, [-1] + self.tensor_shapes[2][1:])
with ops.get_default_graph()._attr_scope({
'_tftrt_op_max_batch_size':
attr_value_pb2.AttrValue(i=self.max_batch_sizes[2])
}):
tensor = tensor * 4.0
tensor = array_ops.reshape(tensor, [-1] + self.tensor_shapes[3][1:])
with ops.get_default_graph()._attr_scope({
'_tftrt_op_max_batch_size':
attr_value_pb2.AttrValue(i=self.max_batch_sizes[3])
}):
tensor += tensor + 5.0
return array_ops.identity(tensor, name='output_0')
class StaticBatchSizeTest(MaxBatchSizesTestBase):
def GraphFn(self, inp):
"""Builds a tf.Graph for the test."""
tensor = inp * 2.0
tensor = array_ops.reshape(tensor, self.tensor_shapes[1])
tensor = tensor + 3.0
tensor = array_ops.reshape(tensor, self.tensor_shapes[2])
tensor = tensor * 4.0
tensor = array_ops.reshape(tensor, self.tensor_shapes[3])
tensor += tensor + 5.0
return array_ops.identity(tensor, name='output_0')
if __name__ == '__main__':
test.main()
@@ -0,0 +1,319 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Basic tests for TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt import utils as trt_utils
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
class SimpleSingleEngineTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing single segment."""
dtype = inp.dtype
conv_filter = constant_op.constant(
[[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]],
name="weights",
dtype=dtype)
conv = nn.conv2d(
input=inp,
filter=conv_filter,
strides=[1, 2, 2, 1],
padding="SAME",
name="conv")
bias = constant_op.constant([4., 1.5, 2., 3., 5., 7.],
name="bias",
dtype=dtype)
added = nn.bias_add(conv, bias, name="bias_add")
relu = nn.relu(added, "relu")
identity = array_ops.identity(relu, "identity")
pool = nn_ops.max_pool(
identity, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
return array_ops.squeeze(pool, name="output_0")
def GetParams(self):
# TODO(aaroey): test graph with different dtypes.
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
[[100, 6, 6, 6]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": [
"weights", "conv", "bias", "bias_add", "relu", "identity",
"max_pool"
]
}
class SimpleMultiEnginesTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing multiple segment."""
dtype = inp.dtype
conv_filter = constant_op.constant(
[[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]],
name="weights",
dtype=dtype)
conv = nn.conv2d(
input=inp,
filter=conv_filter,
strides=[1, 2, 2, 1],
padding="SAME",
name="conv")
c1 = constant_op.constant(
np.random.randn(12, 12, 6), dtype=dtype, name="c1")
p = math_ops.mul(conv, c1, name="mul")
c2 = constant_op.constant(
np.random.randn(12, 12, 6), dtype=dtype, name="c2")
q = math_ops.div(conv, c2, name="div")
edge = self.trt_incompatible_op(q, name="incompatible")
one = constant_op.constant(1, name="one", dtype=dtype)
edge = math_ops.sub(one, edge, name="one_sub")
edge = math_ops.div(edge, edge, name="div1")
r = math_ops.add(edge, edge, name="add")
p = math_ops.sub(p, edge, name="sub")
q = math_ops.mul(q, edge, name="mul1")
s = math_ops.add(p, q, name="add1")
s = math_ops.sub(s, r, name="sub1")
return array_ops.squeeze(s, name="output_0")
def GetParams(self):
# TODO(aaroey): test graph with different dtypes.
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
[[100, 12, 12, 6]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": [
"add", "add1", "c1", "div1", "mul", "mul1", "sub", "sub1", "one",
"one_sub"
],
"TRTEngineOp_001": ["c2", "conv", "div", "weights"]
}
def setUp(self):
super().setUp()
# Disable layout optimizer, since it will convert BiasAdd with NHWC
# format to NCHW format under four dimensional input.
self.DisableNonTrtOptimizers()
def ShouldRunTest(self, run_params):
return (
trt_utils.is_linked_tensorrt_version_greater_equal(8),
"Test is non-hermetic with TensorRT 7",
)
class SimpleMultiEnginesTest2(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing two segments."""
n = inp
for i in range(2):
c = constant_op.constant(float(i), name="c%d" % i)
n = math_ops.add(n, c, name="add%d" % i)
n = math_ops.mul(n, n, name="mul%d" % i)
n = self.trt_incompatible_op(n, name="incompatible")
c = constant_op.constant(2.0, name="c2")
n = math_ops.add(n, c, name="add2")
n = math_ops.mul(n, n, name="mul2")
c = constant_op.constant(3.0, name="c3")
n = math_ops.add(n, c, name="add3")
n = math_ops.mul(n, n, name="mul3")
return array_ops.squeeze(n, name="output_0")
def GetParams(self):
shapes = [[2, 32, 32, 3]]
return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
output_shapes=shapes)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": ["c0", "c1", "add0", "add1", "mul0", "mul1"],
"TRTEngineOp_001": ["c2", "c3", "add2", "add3", "mul2", "mul3"]
}
def ShouldRunTest(self, run_params):
"""Whether to run the test."""
# Disable the test in fp16 mode since multiple matmul and add ops together
# can cause overflow.
return (
(run_params.precision_mode != "FP16") and
not (trt_test.IsQuantizationMode(run_params.precision_mode) and
not run_params.use_calibration)), "test FP32 and non-calibration"
class ConstInputTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing multiple segment."""
n = inp
c = constant_op.constant(1.0, name="c")
# Adds data dependency from the constant op to a trt incompatible op,
# and adds data dependency from the trt incompatible op to the other
# ops, to make sure the constant op cannot be contracted with any trt
# segment that depends on it.
n = self.trt_incompatible_binary_op(n, c, name="incompatible")
n = math_ops.add(n, c, name="add")
n = math_ops.mul(n, n, name="mul")
n = math_ops.add(n, n, name="add1")
n = self.trt_incompatible_op(n, name="incompatible1")
n = math_ops.add(n, c, name="add2")
n = math_ops.mul(n, n, name="mul1")
n = math_ops.add(n, n, name="add3")
return array_ops.squeeze(n, name="output_0")
def GetParams(self):
shapes = [[2, 32, 32, 3]]
return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
output_shapes=shapes)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": ["add", "add1", "mul"],
"TRTEngineOp_001": ["add2", "add3", "mul1"]
}
def ExpectedConnections(self, run_params):
"""Returns the expected edges."""
return {
"input_0": set(),
"c": set(),
"incompatible": {"input_0", "c"},
"TRTEngineOp_000": {"incompatible"},
"incompatible1": {"TRTEngineOp_000"},
"TRTEngineOp_001": {"incompatible1"},
"output_0": {"TRTEngineOp_001"},
}
class ConstDataInputSingleEngineTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing single segment."""
n = inp
c = constant_op.constant(1.0, name="c")
n = math_ops.add(n, c, name="add")
n = math_ops.mul(n, n, name="mul")
n = math_ops.add(n, n, name="add1")
return array_ops.squeeze(n, name="output_0")
def GetParams(self):
shapes = [[2, 32, 32, 3]]
return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
output_shapes=shapes)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {"TRTEngineOp_000": ["c", "add", "add1", "mul"]}
class ConstDataInputMultipleEnginesTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing multiple segment."""
n = inp
c = constant_op.constant(1.0, name="c")
n = math_ops.add(n, c, name="add")
n = math_ops.mul(n, n, name="mul")
n = math_ops.add(n, n, name="add1")
n = self.trt_incompatible_op(n, name="incompatible1")
n = math_ops.add(n, c, name="add2")
n = math_ops.mul(n, n, name="mul1")
n = math_ops.add(n, n, name="add3")
return array_ops.squeeze(n, name="output_0")
def GetParams(self):
shapes = [[2, 32, 32, 3]]
return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
output_shapes=shapes)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": ["add2", "add3", "mul1"],
# Why segment ["add", "add1", "mul"] was assigned segment id 1
# instead of 0: the parent node of this segment is actually const
# node 'c', but it's removed later since it's const output of the
# segment which is not allowed.
"TRTEngineOp_001": ["add", "add1", "mul"]
}
class ControlDependencyTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
"""Create a graph containing multiple segments."""
c1 = constant_op.constant(1.0, name="c1")
c2 = constant_op.constant(2.0, name="c2")
d1 = self.trt_incompatible_op(inp, name="d1")
d2 = self.trt_incompatible_binary_op(inp, inp, name="d2")
with ops.control_dependencies([d1]):
add = math_ops.add(inp, c1, name="add")
mul = math_ops.mul(add, add, name="mul")
add1 = math_ops.add(mul, mul, name="add1")
edge = self.trt_incompatible_op(add1, name="incompatible")
with ops.control_dependencies([d1, d2, add1]):
add2 = math_ops.add(edge, c2, name="add2")
mul1 = math_ops.mul(add2, add2, name="mul1")
add3 = math_ops.add(mul1, mul1, name="add3")
inc1 = self.trt_incompatible_binary_op(d1, add3, name="incompatible1")
inc2 = self.trt_incompatible_binary_op(d2, inc1, name="incompatible2")
return array_ops.squeeze(inc2, name="output_0")
def GetParams(self):
shapes = [[2, 32, 32, 3]]
return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
output_shapes=shapes)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": ["c1", "add", "add1", "mul"],
"TRTEngineOp_001": ["c2", "add2", "add3", "mul1"]
}
def ExpectedConnections(self, run_params):
"""Returns the expected edges."""
return {
"input_0": set(),
"d1": {"input_0"},
"d2": {"input_0"},
"TRTEngineOp_000": {"input_0", "^d1"},
"incompatible": {"TRTEngineOp_000"},
"TRTEngineOp_001": {"incompatible", "^d2"},
"incompatible1": {"TRTEngineOp_001", "d1"},
"incompatible2": {"incompatible1", "d2"},
"output_0": {"incompatible2"},
}
if __name__ == "__main__":
test.main()
@@ -0,0 +1,115 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import unittest
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class BatchMatMultTestBase(trt_test.TfTrtIntegrationTestBase):
"""Base class for BatchMatMult tests."""
# Shape inference of BatchMatMultV2 doesn't work. Use static batch size.
def BuildParams(self, graph_fn, dtype, input_shapes, output_shapes):
return self.BuildParamsWithMask(
graph_fn=graph_fn,
dtype=dtype,
input_shapes=input_shapes,
output_shapes=output_shapes,
input_mask=[[True] * len(s) for s in input_shapes],
output_mask=[[True] * len(s) for s in output_shapes],
extra_inputs=[],
extra_outputs=[])
@classmethod
def setUpClass(cls):
if cls is BatchMatMultTestBase:
raise unittest.SkipTest(
"BatchMatMultTestBase defines base class for other test.")
super(BatchMatMultTestBase, cls).setUpClass()
class BatchMatMulTwoTensorTest(BatchMatMultTestBase):
"""Testing conversion of BatchMatMul where both inputs are tensors."""
def GraphFn(self, inp, inp1):
x1 = math_ops.matmul(inp, inp1, name="matmul")
# Relu to reach minimum segment size.
x1 = nn.relu(x1, name="relu")
return array_ops.identity(x1, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32,
[[12, 5, 8, 12], [12, 5, 12, 7]], [[12, 5, 8, 7]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {"TRTEngineOp_000": ["matmul", "relu"]}
class BatchMatMulWeightBroadcastTest(BatchMatMultTestBase):
"""Testing BatchMatMulV2: one operand is weight and both have same rank."""
def ShouldAllowTF32Computation(self):
return False
def GraphFn(self, inp):
dtype = inp.dtype
b = constant_op.constant(
np.random.randn(1, 5, 7), dtype=dtype, name="kernel")
x1 = math_ops.matmul(inp, b, name="matmul")
return array_ops.identity(x1, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[12, 9, 5]],
[[12, 9, 7]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {"TRTEngineOp_000": ["matmul", "kernel"]}
class BatchMatMulWeightBroadcastDims2Test(BatchMatMultTestBase):
"""Testing BatchMatMulV2: weight operand must be broadcasted."""
def ShouldAllowTF32Computation(self):
return False
def GraphFn(self, inp):
dtype = inp.dtype
b = constant_op.constant(np.random.randn(5, 7), dtype=dtype, name="kernel")
x1 = math_ops.matmul(inp, b, name="matmul")
return array_ops.identity(x1, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[12, 9, 5]],
[[12, 9, 7]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {"TRTEngineOp_000": ["matmul", "kernel"]}
if __name__ == "__main__":
test.main()
@@ -0,0 +1,130 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class BiasaddMatMulTest(trt_test.TfTrtIntegrationTestBase):
"""Testing conversion of BiasAdd MatMul in TF-TRT conversion."""
def _ConstOp(self, shape):
return constant_op.constant(np.random.randn(*shape), dtype=dtypes.float32)
def GraphFn(self, x):
input_matrix_rows = 4
input_matrix_columns = 144
b = self._ConstOp((input_matrix_columns, 4))
x1 = math_ops.matmul(x, b)
b = self._ConstOp((1, 4))
x1 = x1 + b
b = self._ConstOp((input_matrix_rows, 144))
x2 = self.trt_incompatible_op(x)
x2 = math_ops.matmul(x2, b, transpose_a=True)
x2 = gen_array_ops.reshape(x2, [4, -1])
x2 = self.trt_incompatible_op(x2)
b = self._ConstOp((4, input_matrix_columns))
x3 = math_ops.matmul(x, b, transpose_b=True)
b = self._ConstOp((16, input_matrix_rows))
x4 = self.trt_incompatible_op(x)
x4 = math_ops.matmul(x4, b, transpose_b=True, transpose_a=True)
x4 = gen_array_ops.reshape(x4, [4, -1])
x4 = self.trt_incompatible_op(x4)
# Note that tf.nn.bias_add supports up to 5 dimensions.
b = self._ConstOp((input_matrix_columns, 48))
x5 = math_ops.matmul(x, b)
b = self._ConstOp((48,))
x5 = nn.bias_add(x5, b)
# TODO(b/154672994): Put the reshape back when the bug is fixed.
# x5 = gen_array_ops.reshape(x5, [4, -1])
x6 = gen_array_ops.reshape(x, [4, 24, 6])
b = self._ConstOp((6,))
x6 = nn.bias_add(x6, b, data_format="NHWC")
x6 = gen_array_ops.reshape(x6, [4, -1])
x7 = gen_array_ops.reshape(x, [4, 12, 4, 3])
b = self._ConstOp((3,))
x7 = nn.bias_add(x7, b, data_format="NHWC")
x7 = gen_array_ops.reshape(x7, [4, -1])
x8 = gen_array_ops.reshape(x, [4, 4, 3, 2, 6])
b = self._ConstOp((6,))
x8 = nn.bias_add(x8, b, data_format="NHWC")
x8 = gen_array_ops.reshape(x8, [4, -1])
x9 = gen_array_ops.reshape(x, [4, 12, 3, 2, 2])
b = self._ConstOp((12,))
x9 = nn.bias_add(x9, b, data_format="NCHW")
x9 = gen_array_ops.reshape(x9, [4, -1])
x10 = gen_array_ops.reshape(x, [4, 3, 4, 12])
b = self._ConstOp((3,))
x10 = nn.bias_add(x10, b, data_format="NCHW")
x10 = gen_array_ops.reshape(x10, [4, -1])
x11 = gen_array_ops.reshape(x, [4, 6, 24])
b = self._ConstOp((6,))
x11 = nn.bias_add(x11, b, data_format="NCHW")
x11 = gen_array_ops.reshape(x11, [4, -1])
out = array_ops.concat([x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11],
axis=-1)
return array_ops.squeeze(out, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[4, 144]],
[[4, 6680]])
def setUp(self):
super().setUp()
# Disable layout optimizer, since it will convert BiasAdd with NHWC
# format to NCHW format under four dimensional input.
self.DisableNonTrtOptimizers()
def GetMaxBatchSize(self, run_params):
"""Returns the max_batch_size that the converter should use for tests."""
if run_params.dynamic_engine:
return None
return 4
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
if run_params.dynamic_shape:
# Increased conversion rate in dynamic shape mode due to a few additional
# conversions for MatMul, Reshape and Concat ops. This increases the size
# of the candidate segments and results in two more TrtEngineOps.
return ["TRTEngineOp_000", "TRTEngineOp_001", "TRTEngineOp_002"]
else:
return ["TRTEngineOp_000"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,80 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import os
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
class BinaryTensorWeightBroadcastTest(trt_test.TfTrtIntegrationTestBase):
"""Tests for scale & elementwise layers in TF-TRT."""
def _ConstOp(self, shape):
return constant_op.constant(np.random.randn(*shape), dtype=dtypes.float32)
def GraphFn(self, x):
for weights_shape in [
(1,), # scale
(24, 1, 1), # scale
(24, 24, 20), # scale
(20,), # elementwise
(1, 24, 1, 1), # elementwise
(1, 24, 24, 1), # elementwise
(1, 24, 24, 20), # elementwise
(24, 20), # elementwise
]:
a = self._ConstOp(weights_shape)
f = x + a
x = self.trt_incompatible_op(f)
a = self._ConstOp(weights_shape)
f = a + x
x = self.trt_incompatible_op(f)
return gen_array_ops.reshape(x, [5, -1], name="output_0")
def GetParams(self):
# TODO(aaroey): test graph with different dtypes.
return self.BuildParams(self.GraphFn, dtypes.float32, [[10, 24, 24, 20]],
[[5, 23040]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
# The final reshape op is converted only in dynamic shape mode. This op is
# placed into a new engine due to the preceding trt_incompatible_ops.
num_engines = 17 if run_params.dynamic_shape else 16
return [f"TRTEngineOp_{i:03d}" for i in range(num_engines)]
# TODO(b/176540862): remove this routine to disallow native segment execution
# for TensorRT 7+.
def setUp(self):
super().setUp()
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
gpus = config.list_physical_devices("GPU")
logging.info("Found the following GPUs:")
for gpu in gpus:
logging.info(f"\t- {gpu}")
config.set_memory_growth(gpu, True)
if __name__ == "__main__":
test.main()
@@ -0,0 +1,51 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.compiler.tensorrt import utils as trt_utils
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class BoolTest(trt_test.TfTrtIntegrationTestBase):
"""Test for boolean operations in TF-TRT."""
def GraphFn(self, x1, x2):
x = math_ops.logical_and(x1, x2)
x = math_ops.logical_or(x, x2)
q = math_ops.not_equal(x, x2)
q = math_ops.logical_not(q)
return array_ops.identity(q, name="output_0")
def GetParams(self):
shape = [2, 32, 32, 3]
return self.BuildParams(self.GraphFn, dtypes.bool, [shape, shape], [shape])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
reason = "Boolean ops are not implemented "
return (run_params.dynamic_shape, reason + "in ImplicitBatch mode") \
if trt_utils.is_linked_tensorrt_version_greater_equal(8, 2, 0) \
else (False, reason + "for TRT < 8.2.0")
if __name__ == "__main__":
test.main()
@@ -0,0 +1,60 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Test conversion of graphs involving INT32 tensors and operations."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class CastInt32ToFp32Test(trt_test.TfTrtIntegrationTestBase):
"""Tests cast to FP32 are split in FP16 mode."""
def _ConstOp(self, shape, dtype):
return constant_op.constant(np.random.randn(*shape), dtype=dtype)
def GraphFn(self, x):
b_f = self._ConstOp((1, 10), dtypes.float16)
x_f = math_ops.cast(x, dtypes.float16)
x_f = math_ops.mul(x_f, b_f) # FP16 Multiply
x_f = math_ops.cast(x_f, dtypes.float32)
b_f = self._ConstOp((1, 10), dtypes.float32)
x_f = math_ops.add(x_f, b_f) # FP32 Add
return array_ops.identity(x_f, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[1, 10]], [[1, 10]])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
return {"TRTEngineOp_000": ["AddV2", "Cast", "Const", "Mul"]}
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-02
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-02
if __name__ == "__main__":
test.main()
@@ -0,0 +1,222 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Script to test TF-TensorRT conversion of CombinedNMS op."""
import os
from tensorflow.python.compiler.tensorrt import utils as trt_utils
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import image_ops_impl
from tensorflow.python.platform import test
class CombinedNmsTest(trt_test.TfTrtIntegrationTestBase):
"""Test for CombinedNMS op in TF-TRT."""
def setUp(self):
super().setUp()
self.num_boxes = 200
def GraphFn(self, boxes, scores):
max_total_size = 3
score_threshold = 0.1
iou_threshold = 0.5
# Shapes
max_total_size_tensor = constant_op.constant(
max_total_size, dtype=dtypes.int32, name='max_total_size')
iou_threshold_tensor = constant_op.constant(
iou_threshold, dtype=dtypes.float32, name='iou_threshold')
score_threshold_tensor = constant_op.constant(
score_threshold, dtype=dtypes.float32, name='score_threshold')
nms_output = image_ops_impl.combined_non_max_suppression(
boxes,
scores,
max_total_size_tensor,
max_total_size_tensor,
iou_threshold_tensor,
score_threshold_tensor,
name='combined_nms')
return [
array_ops.identity(output, name=('output_%d' % i))
for i, output in enumerate(nms_output)
]
def GetParams(self):
# Parameters
q = 1
batch_size = 2
num_classes = 2
max_total_size = 3
boxes_shape = [batch_size, self.num_boxes, q, 4]
scores_shape = [batch_size, self.num_boxes, num_classes]
nmsed_boxes_shape = [batch_size, max_total_size, 4]
nmsed_scores_shape = [batch_size, max_total_size]
nmsed_classes_shape = [batch_size, max_total_size]
valid_detections_shape = [batch_size]
return self.BuildParams(self.GraphFn, dtypes.float32,
[boxes_shape, scores_shape], [
nmsed_boxes_shape, nmsed_scores_shape,
nmsed_classes_shape, valid_detections_shape
])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
if not run_params.dynamic_shape:
return {
'TRTEngineOp_000': [
'combined_nms/CombinedNonMaxSuppression', 'max_total_size',
'iou_threshold', 'score_threshold'
]
}
else:
# The CombinedNMS op is currently not converted in dynamic shape mode.
# This branch shall be removed once the converter is updated to handle
# input with dynamic shape.
return dict()
def ShouldRunTest(self, run_params):
should_run, reason = super().ShouldRunTest(run_params)
should_run = should_run and \
not trt_test.IsQuantizationMode(run_params.precision_mode)
reason += ' and precision != INT8'
# Only run for TRT 7.1.3 and above.
return should_run and trt_utils.is_linked_tensorrt_version_greater_equal(
7, 1, 3), reason + ' and >= TRT 7.1.3'
class CombinedNmsExecuteNativeSegmentTest(CombinedNmsTest):
def setUp(self):
super().setUp()
os.environ['TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION'] = 'True'
def tearDown(self):
super().tearDown()
os.environ['TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION'] = 'False'
def GetMaxBatchSize(self, run_params):
"""Returns the max_batch_size that the converter should use for tests."""
if run_params.dynamic_engine:
return None
# Build the engine with the allowed max_batch_size less than the actual
# max_batch_size, to fore the runtime to execute the native segment. This
# is to test that combined_non_max_suppression, which doesn't have a TF GPU
# implementation, can be executed natively even though the it is in the
# the graph for the TRTEngineOp with a GPU as a default device.
return super().GetMaxBatchSize(run_params) - 1
def ShouldRunTest(self, run_params):
should_run, reason = super().ShouldRunTest(run_params)
# max_batch_size is only useful for selecting static engines. As such,
# we shouldn't run the test for dynamic engines.
return should_run and \
not run_params.dynamic_engine, reason + ' and static engines'
class CombinedNmsTestTopK(CombinedNmsTest):
"""Test for CombinedNMS TopK op in TF-TRT."""
def GraphFn(self, pre_nms_boxes, pre_nms_scores, max_boxes_to_draw,
max_detetion_points):
iou_threshold = 0.1
score_threshold = 0.001
max_output_size_per_class_tensor = constant_op.constant(
max_detetion_points,
dtype=dtypes.int32,
name='max_output_size_per_class')
max_total_size_tensor = constant_op.constant(
max_boxes_to_draw, dtype=dtypes.int32, name='max_total_size')
iou_threshold_tensor = constant_op.constant(
iou_threshold, dtype=dtypes.float32, name='iou_threshold')
score_threshold_tensor = constant_op.constant(
score_threshold, dtype=dtypes.float32, name='score_threshold')
nms_output = image_ops_impl.combined_non_max_suppression(
pre_nms_boxes,
pre_nms_scores,
max_output_size_per_class=max_output_size_per_class_tensor,
max_total_size=max_total_size_tensor,
iou_threshold=iou_threshold_tensor,
score_threshold=score_threshold_tensor,
pad_per_class=False,
name='combined_nms')
return [
array_ops.identity(output, name=('output_%d' % i))
for i, output in enumerate(nms_output)
]
def GetParams(self):
# Parameters
batch_size = 1
max_detetion_points = 2048
num_classes = 90
max_boxes_to_draw = 30
# Inputs
pre_nms_boxes_shape = [batch_size, max_detetion_points, 1, 4]
pre_nms_scores_shape = [batch_size, max_detetion_points, num_classes]
# Outputs
nmsed_boxes_shape = [batch_size, max_boxes_to_draw, 4]
nmsed_scores_shape = [batch_size, max_boxes_to_draw]
nmsed_classes_shape = [batch_size, max_boxes_to_draw]
valid_detections_shape = [batch_size]
def _get_graph_fn(x, y):
return self.GraphFn(
x,
y,
max_boxes_to_draw=max_boxes_to_draw,
max_detetion_points=max_detetion_points)
return self.BuildParams(_get_graph_fn, dtypes.float32,
[pre_nms_boxes_shape, pre_nms_scores_shape], [
nmsed_boxes_shape, nmsed_scores_shape,
nmsed_classes_shape, valid_detections_shape
])
class CombinedNmsTopKOverride(CombinedNmsTest):
def setUp(self):
super().setUp()
self.num_boxes = 5000
os.environ['TF_TRT_ALLOW_NMS_TOPK_OVERRIDE'] = '1'
def tearDown(self):
super().tearDown()
os.environ['TF_TRT_ALLOW_NMS_TOPK_OVERRIDE'] = '0'
def GetMaxBatchSize(self, run_params):
"""Returns the max_batch_size that the converter should use for tests."""
if run_params.dynamic_engine:
return None
return super().GetMaxBatchSize(run_params)
if __name__ == '__main__':
test.main()
@@ -0,0 +1,73 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.platform import test
class ConcatenationTest(trt_test.TfTrtIntegrationTestBase):
"""Testing Concatenation in TF-TRT conversion."""
def GraphFn(self, x):
dtype = x.dtype
# scale
a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
r1 = x / a
a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
r2 = a / x
a = constant_op.constant(np.random.randn(1, 3, 1), dtype=dtype)
r3 = a + x
a = constant_op.constant(np.random.randn(1, 3, 1), dtype=dtype)
r4 = x * a
a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
r5 = x - a
a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
r6 = a - x
a = constant_op.constant(np.random.randn(3, 1), dtype=dtype)
r7 = x - a
a = constant_op.constant(np.random.randn(3, 1), dtype=dtype)
r8 = a - x
a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
r9 = gen_math_ops.maximum(x, a)
a = constant_op.constant(np.random.randn(3, 1), dtype=dtype)
r10 = gen_math_ops.minimum(a, x)
a = constant_op.constant(np.random.randn(3), dtype=dtype)
r11 = x * a
a = constant_op.constant(np.random.randn(1), dtype=dtype)
r12 = a * x
concat1 = array_ops.concat([r1, r2, r3, r4, r5, r6], axis=-1)
concat2 = array_ops.concat([r7, r8, r9, r10, r11, r12], axis=3)
x = array_ops.concat([concat1, concat2], axis=-1)
return gen_array_ops.reshape(x, [2, -1], name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 3, 3, 1]],
[[2, 126]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,61 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Script to test TF-TensorRT integration."""
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class ConstBroadcastTest(trt_test.TfTrtIntegrationTestBase):
"""Test for Constant broadcasting in TF-TRT."""
def GraphFn(self, x):
"""Return the expected graph to convert."""
dtype = x.dtype
filt1 = constant_op.constant(
0.3, shape=(3, 3, 2, 1), dtype=dtype, name='filt1')
y1 = nn.conv2d(x, filt1, strides=[1, 1, 1, 1], padding='SAME', name='y1')
z1 = nn.relu(y1, name='z1')
filt2 = constant_op.constant(
0.3, shape=(3, 3, 1, 1), dtype=dtype, name='filt2')
y2 = nn.conv2d(z1, filt2, strides=[1, 1, 1, 1], padding='SAME', name='y2')
z2 = nn.relu(y2, name='z')
filt3 = constant_op.constant(
0.3, shape=(3, 3, 1, 1), dtype=dtype, name='filt3')
y3 = nn.conv2d(z2, filt3, strides=[1, 1, 1, 1], padding='SAME', name='y3')
return nn.relu(y3, name='output_0')
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[5, 12, 12, 2]],
[[5, 12, 12, 1]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ['TRTEngineOp_000']
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 1.e-04 if run_params.precision_mode == 'FP32' else 1.e-02
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 1.e-04 if run_params.precision_mode == 'FP32' else 1.e-02
if __name__ == '__main__':
test.main()
@@ -0,0 +1,207 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
def conv2d_layer(inputs,
filters,
kernel_size,
strides=(1, 1),
padding="valid",
data_format="channels_last",
dilation_rate=(1, 1)):
dtype = inputs.dtype
c_axis = -1 if data_format == "channels_last" else 1
nchan = inputs.shape[c_axis]
weights_shape = (kernel_size[0], kernel_size[1], nchan, filters)
weights = constant_op.constant(np.random.randn(*weights_shape), dtype=dtype)
padding = padding.upper()
if data_format == "channels_last":
strides = [1] + list(strides) + [1]
dilations = [1] + list(dilation_rate) + [1]
data_format = "NHWC"
else:
strides = [1, 1] + list(strides)
dilations = [1, 1] + list(dilation_rate)
data_format = "NCHW"
return gen_nn_ops.conv2d(
inputs,
weights,
strides=strides,
padding=padding,
dilations=dilations,
data_format=data_format)
def div_round_up(n, d):
return (n - 1) // d + 1
def build_graph(inp,
dtype,
num_filters,
data_format,
kernel_sizes,
dilation_rates,
padding="same"):
results = []
for kernel_size in kernel_sizes:
for dilation_rate in dilation_rates:
result = conv2d_layer(inp, num_filters, kernel_size, (1, 1), padding,
data_format, dilation_rate)
results.append(result)
output = sum(results)
return array_ops.identity(output, name="output_0")
class Conv2DNCHWTest(trt_test.TfTrtIntegrationTestBase):
"""Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion."""
def GraphFn(self, inp):
np.random.seed(1234)
return build_graph(
inp=inp,
dtype=dtypes.float32,
num_filters=5,
data_format="channels_first",
kernel_sizes=[(3, 3), (3, 2)],
dilation_rates=[(1, 1), (2, 3)])
def GetParams(self):
# TODO(aaroey): test graph with different dtypes.
return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 3, 7, 11]],
[[13, 5, 7, 11]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
if trt_test.IsQuantizationWithCalibration(run_params):
return 4e-02
return super(Conv2DNCHWTest, self).ExpectedAbsoluteTolerance(run_params)
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
if trt_test.IsQuantizationWithCalibration(run_params):
return 4e-02
return super(Conv2DNCHWTest, self).ExpectedRelativeTolerance(run_params)
class Conv2DNHWCTest(trt_test.TfTrtIntegrationTestBase):
"""Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion."""
def GraphFn(self, inp):
np.random.seed(1234)
return build_graph(
inp=inp,
dtype=dtypes.float32,
num_filters=5,
data_format="channels_last",
kernel_sizes=[(3, 3), (3, 2)],
dilation_rates=[(1, 1), (2, 3)])
def GetParams(self):
# TODO(aaroey): test graph with different dtypes.
return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 7, 11, 3]],
[[13, 7, 11, 5]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
class Conv2DStridedNCHWTest(trt_test.TfTrtIntegrationTestBase):
"""Testing conversion of strided Conv2D (data_format=NCHW)."""
def GraphFn(self, inp):
np.random.seed(1234)
num_filters = 5
output = inp
output = conv2d_layer(
output,
num_filters, (3, 2),
strides=(2, 2),
padding="same",
data_format="channels_first")
output = conv2d_layer(
output,
num_filters, (3, 3),
strides=(2, 2),
dilation_rate=(2, 3),
padding="same",
data_format="channels_first")
return array_ops.identity(output, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 3, 7, 11]],
[[13, 5, 2, 3]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 5.e-01 if run_params.precision_mode == "INT8" else 1.e-02
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 5.e-00 if run_params.precision_mode == "INT8" else 1.e-02
class Conv2DTranposeTest(trt_test.TfTrtIntegrationTestBase):
"""Testing conversion of conv2d_transpose (AKA Conv2DBackpropInput)"""
def GraphFn(self, inp):
np.random.seed(1234)
dtype = inp.dtype
n, c, h, w = 13, 3, 7, 11
num_filters = 8
weights_shape = [2, 2, num_filters, c]
weights = constant_op.constant(np.random.randn(*weights_shape), dtype=dtype)
output_shape = constant_op.constant([n, num_filters, h * 2, w * 2],
dtype=dtypes.int32)
output = nn_ops.conv2d_transpose(
inp,
weights,
output_shape,
strides=[1, 1, 2, 2],
padding="SAME",
data_format="NCHW")
return array_ops.identity(output, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 3, 7, 11]],
[[13, 8, 14, 22]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,74 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration with data dependent shapes"""
import os
from unittest import SkipTest # pylint: disable=g-importing-member
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class TrtModeTestBase(trt_test.TfTrtIntegrationTestBase):
"""Creates a network that has data dependent shapes."""
def GraphFn(self, x):
# With the unique() op we create a tensor with data dependent shape.
x = math_ops.floor(x * 10)
y, idx = array_ops.unique(x)
y = y * 2 + y
# The rest is only needed to ensure that the output has the same size as
# the input (expected by the test harness).
padding = array_ops.constant([0])
n = array_ops.shape(x) - array_ops.shape(y)
padding = array_ops.concat([padding, n], 0)
padding = array_ops.expand_dims(padding, 0)
y = array_ops.pad(y, padding)
return array_ops.identity(y, name="output_0")
def ShouldRunTest(self, run_params):
# We have a single dimension, and that is changed, therefore the graph can
# only be converted in dynamic shape mode.
return (run_params.dynamic_engine and run_params.is_v2 and
run_params.dynamic_shape and
run_params.use_calibration, "test v2 dynamic engine and "
"calibration")
def setUp(self):
super().setUp()
# The input shape depends on random input data. It can happen that we do
# not have engine for the actual shape. Therefore we enable native segment
# execution.
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
def tearDown(self):
super().tearDown()
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "False"
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[12]], [[12]])
def ExpectedEnginesToBuild(self, run_params):
return ["TRTEngineOp_000", "TRTEngineOp_001"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,102 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Script to test TF-TRT INT8 conversion without calibration on Mnist model."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class DynamicInputShapesTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, x):
conv_filter1 = constant_op.constant(
np.ones([3, 3, 1, 8]), name="weights1", dtype=dtypes.float32)
bias1 = constant_op.constant(np.random.randn(8), dtype=dtypes.float32)
x = nn.conv2d(
input=x,
filter=conv_filter1,
strides=[1, 1, 1, 1],
padding="SAME",
name="conv")
x = nn.bias_add(x, bias1)
x = nn.relu(x)
conv_filter2 = constant_op.constant(
np.ones([3, 3, 8, 1]), name="weights2", dtype=dtypes.float32)
bias2 = constant_op.constant(np.random.randn(1), dtype=dtypes.float32)
x = nn.conv2d(
input=x,
filter=conv_filter2,
strides=[1, 1, 1, 1],
padding="SAME",
name="conv")
x = nn.bias_add(x, bias2)
return array_ops.identity(x, name="output")
def GetParams(self):
# TODO(laigd): we should test the following cases:
# - batch size is not changed, other dims are changing
# - batch size is decreasing, other dims are identical
# - batch size is decreasing, other dims are changing
# - batch size is increasing, other dims are identical
# - batch size is increasing, other dims are changing
input_dims = [[[1, 5, 5, 1]], [[10, 5, 5, 1]], [[3, 5, 5, 1]],
[[1, 5, 5, 1]], [[1, 3, 1, 1]], [[2, 9, 9, 1]],
[[1, 224, 224, 1]], [[1, 128, 224, 1]]]
expected_output_dims = input_dims
return trt_test.TfTrtIntegrationTestParams(
graph_fn=self.GraphFn,
input_specs=[
tensor_spec.TensorSpec([None, None, None, 1], dtypes.float32,
"input")
],
output_specs=[
tensor_spec.TensorSpec([None, None, None, 1], dtypes.float32,
"output")
],
input_dims=input_dims,
expected_output_dims=expected_output_dims)
def setUp(self):
super().setUp()
# Disable layout optimizer, since it will convert BiasAdd with NHWC
# format to NCHW format under four dimensional input.
self.DisableNonTrtOptimizers()
def ExpectedEnginesToBuild(self, run_params):
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
return (run_params.dynamic_engine and not trt_test.IsQuantizationMode(
run_params.precision_mode)), "test dynamic engine and non-INT8"
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-01
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-01
if __name__ == "__main__":
test.main()
@@ -0,0 +1,52 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""This test checks a situation where the same tensor is considered as an output
multiple times because it has been duplicated by 2+ identity ops. Previously,
the tensor would be renamed multiple times, overwriting the output binding name
which resulted in a runtime error when the binding would not be found.
"""
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class IdentityTest(trt_test.TfTrtIntegrationTestBase):
"""Testing engine with the same tensor repeated as output via identity."""
def GraphFn(self, x):
x1 = math_ops.exp(x)
x1 = x1 + x
out1 = array_ops.identity(x1, name='output_0')
out2 = array_ops.identity(x1, name='output_1')
iden1 = array_ops.identity(x1)
out3 = array_ops.identity(iden1, name='output_2')
return [out1, out2, out3]
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 32]],
[[100, 32]] * 3)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ['TRTEngineOp_000']
if __name__ == '__main__':
test.main()
@@ -0,0 +1,86 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Test conversion of graphs involving INT32 tensors and operations."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class ExcludeUnsupportedInt32Test(trt_test.TfTrtIntegrationTestBase):
"""Test exclusion of ops which are not supported in INT32 mode by TF-TRT"""
def _ConstOp(self, shape, dtype):
return constant_op.constant(np.random.randn(*shape), dtype=dtype)
def GraphFn(self, x):
dtype = x.dtype
b = self._ConstOp((4, 10), dtype)
x = math_ops.matmul(x, b)
b = self._ConstOp((10,), dtype)
x = nn.bias_add(x, b)
return array_ops.identity(x, name='output_0')
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.int32, [[100, 4]], [[100, 10]])
def setUp(self):
super().setUp()
# Disable layout optimizer, since it will convert BiasAdd with NHWC
# format to NCHW format under four dimensional input.
self.DisableNonTrtOptimizers()
def GetMaxBatchSize(self, run_params):
"""Returns the max_batch_size that the converter should use for tests."""
if run_params.dynamic_engine:
return None
return 100
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return []
class CalibrationInt32Support(trt_test.TfTrtIntegrationTestBase):
"""Test execution of calibration with int32 input"""
def GraphFn(self, inp):
# Can use any op that is converted to TRT with int32 inputs
inp_transposed = array_ops.transpose(inp, [0, 3, 2, 1], name='transpose_0')
return array_ops.identity(inp_transposed, name='output_0')
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.int32, [[3, 4, 5, 6]],
[[3, 6, 5, 4]])
def ShouldRunTest(self, run_params):
# Although test passes with all configurations but only
# execute INT8 with use_calibration=True because
# that is the purpose of the test.
return trt_test.IsQuantizationWithCalibration(
run_params), 'test calibration and INT8'
def ExpectedEnginesToBuild(self, run_params):
return ['TRTEngineOp_000']
if __name__ == '__main__':
test.main()
@@ -0,0 +1,66 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Test LRUCache by running different input batch sizes on same network."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class LRUCacheTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, x):
bias = constant_op.constant(
np.random.randn(1, 10, 10, 1), dtype=dtypes.float32)
x = math_ops.add(x, bias)
x = nn.relu(x)
return array_ops.identity(x, name="output")
def GetParams(self):
dtype = dtypes.float32
input_dims = [[[1, 10, 10, 2]], [[2, 10, 10, 2]], [[4, 10, 10, 2]],
[[2, 10, 10, 2]]]
expected_output_dims = [[[1, 10, 10, 2]], [[2, 10, 10, 2]], [[4, 10, 10,
2]],
[[2, 10, 10, 2]]]
return trt_test.TfTrtIntegrationTestParams(
graph_fn=self.GraphFn,
input_specs=[
tensor_spec.TensorSpec([None, 10, 10, 2], dtypes.float32, "input")
],
output_specs=[
tensor_spec.TensorSpec([None, 10, 10, 1], dtypes.float32, "output")
],
input_dims=input_dims,
expected_output_dims=expected_output_dims)
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
return (run_params.dynamic_engine and not trt_test.IsQuantizationMode(
run_params.precision_mode)), "test dynamic engine and non-INT8"
if __name__ == "__main__":
test.main()
@@ -0,0 +1,70 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
@test_util.run_all_without_tensor_float_32("Uses matmul")
class MemoryAlignmentTest(trt_test.TfTrtIntegrationTestBase):
"""Testing conversion of BatchMatMul in TF-TRT conversion."""
def GraphFn(self, inp):
dtype = inp.dtype
e1 = constant_op.constant(
np.random.randn(1, 1, 3, 5), name="kernel_1", dtype=dtype)
e2 = constant_op.constant(
np.random.randn(1, 1, 5, 10), name="kernel_2", dtype=dtype)
conv = nn.conv2d(
input=inp,
filter=e1,
strides=[1, 1, 1, 1],
padding="VALID",
name="conv")
out = nn.conv2d(
input=conv,
filter=e2,
strides=[1, 1, 1, 1],
padding="VALID",
name="conv_2")
return array_ops.squeeze(out, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 15, 15, 3]],
[[2, 15, 15, 10]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 1.e-06 if run_params.precision_mode == "FP32" else 1.e-02
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 0.1
if __name__ == "__main__":
test.main()
@@ -0,0 +1,73 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class MultiConnectionNeighborEngineTest(trt_test.TfTrtIntegrationTestBase):
"""Test for multi connection neighboring nodes wiring tests in TF-TRT."""
def GraphFn(self, x):
dtype = x.dtype
e = constant_op.constant(
np.random.normal(.05, .005, [3, 2, 3, 4]), name="weights", dtype=dtype)
conv = nn.conv2d(
input=x,
filter=e,
data_format="NCHW",
strides=[1, 1, 1, 1],
padding="VALID",
name="conv")
b = constant_op.constant(
np.random.normal(2.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
t = conv + b
b = constant_op.constant(
np.random.normal(5.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
q = conv - b
edge = self.trt_incompatible_op(q)
b = constant_op.constant(
np.random.normal(5.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
d = b + conv
edge3 = self.trt_incompatible_op(d)
edge1 = self.trt_incompatible_op(conv)
t = t - edge1
q = q + edge
t = t + q
t = t + d
t = t - edge3
return array_ops.squeeze(t, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 3, 7, 5]],
[[2, 4, 5, 4]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000", "TRTEngineOp_001"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,62 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
class NeighboringEngineTest(trt_test.TfTrtIntegrationTestBase):
"""Neighboring node wiring tests in TF-TRT conversion."""
def GraphFn(self, x):
dtype = x.dtype
e = constant_op.constant(
np.random.normal(.3, 0.05, [3, 2, 3, 4]), name="weights", dtype=dtype)
conv = nn.conv2d(
input=x,
filter=e,
data_format="NCHW",
strides=[1, 1, 1, 1],
padding="VALID",
name="conv")
b = constant_op.constant(
np.random.normal(1.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
t = math_ops.mul(conv, b, name="mul")
e = self.trt_incompatible_op(conv, name="incompatible")
t = math_ops.sub(t, e, name="sub")
return array_ops.squeeze(t, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 3, 7, 5]],
[[2, 4, 5, 4]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": ["bias", "mul", "sub"],
"TRTEngineOp_001": ["weights", "conv"]
}
if __name__ == "__main__":
test.main()
@@ -0,0 +1,133 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
def _GraphFn(x, add_quantization_nodes):
def _Quantize(x, r):
if add_quantization_nodes:
x = gen_array_ops.fake_quant_with_min_max_vars(x, -r, r)
return x
x = _Quantize(x, 10.0)
x = x + 5
x = _Quantize(x, 15.0)
x = x - 5
x = _Quantize(x, 10.0)
x = x * 0.1
x = _Quantize(x, 1.0)
w = constant_op.constant(np.ones((8, 1)), dtype=dtypes.float32)
x = math_ops.matmul(x, w)
x = _Quantize(x, 10.0)
return array_ops.identity(x, name="output_0")
def _GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[8, 8]], [[8, 1]])
class QuantizationMissingAllRangesTest(trt_test.TfTrtIntegrationTestBase):
"""Create a graph containing single segment with no quantization ranges."""
def GraphFn(self, x):
return _GraphFn(x, add_quantization_nodes=False)
def GetParams(self):
return _GetParams(self)
def ShouldRunTest(self, run_params):
# Only test static engine mode, with or without calibration.
return (trt_test.IsQuantizationMode(run_params.precision_mode) and
not run_params.convert_online and not run_params.dynamic_engine
), "test static engine, offline conversion and INT8"
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
# In static engine mode with calibration, it should build a calibration
# engine.
# In static engine mode without calibration, the engine building will
# succeed but fall back to non-quantized ops.
return ["TRTEngineOp_000"]
class QuantizationWithRangesTest(trt_test.TfTrtIntegrationTestBase):
"""Create a graph containing single segment with no quantization ranges."""
def GraphFn(self, x):
return _GraphFn(x, add_quantization_nodes=True)
def GetParams(self):
return _GetParams(self)
def ShouldRunTest(self, run_params):
# Test static/dynamic engine with/without calibration.
return (trt_test.IsQuantizationMode(run_params.precision_mode) and
not run_params.convert_online), "test offline conversion and INT8"
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
class NonQuantizedPrecisionsWithRangesTest(trt_test.TfTrtIntegrationTestBase):
"""Create a graph containing single segment with no quantization ranges."""
def GraphFn(self, x):
return _GraphFn(x, add_quantization_nodes=True)
def GetParams(self):
return _GetParams(self)
def ShouldRunTest(self, run_params):
# Only test FP32/FP16 mode.
return not trt_test.IsQuantizationMode(
run_params.precision_mode), "test non-INT8"
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
# The fake quant ops are not supported in FP32/FP16 mode, and will split the
# graph into two TRT segments.
return ["TRTEngineOp_000", "TRTEngineOp_001"]
def ExpectedAbsoluteTolerance(self, run_params):
"""The absolute tolerance to compare floating point results."""
return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
def ExpectedRelativeTolerance(self, run_params):
"""The relative tolerance to compare floating point results."""
return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
if __name__ == "__main__":
test.main()
@@ -0,0 +1,85 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class RankTwoTest(trt_test.TfTrtIntegrationTestBase):
"""Test for rank 2 input in TF-TRT."""
def GraphFn(self, x1, x2):
# Two paths: first with rank 2 input, second with rank 4 input.
outputs = []
xs = [x1, x2]
for i in range(2):
x = xs[i]
c = constant_op.constant(1.0, name="c%d_1" % i)
q = math_ops.add(x, c, name="add%d_1" % i)
q = math_ops.abs(q, name="abs%d_1" % i)
c = constant_op.constant(2.2, name="c%d_2" % i)
q = math_ops.add(q, c, name="add%d_2" % i)
q = math_ops.abs(q, name="abs%d_2" % i)
c = constant_op.constant(3.0, name="c%d_3" % i)
q = math_ops.add(q, c, name="add%d_3" % i)
if i == 0:
axis = constant_op.constant(-1, dtype=dtypes.int32, name="axis")
for j in range(2):
q = array_ops.expand_dims(q, axis, name="expand%d_%d" % (i, j))
q = self.trt_incompatible_op(q)
q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i)
outputs.append(q)
# Combine both paths
q = math_ops.add(outputs[0], outputs[1], name="add")
return array_ops.squeeze(q, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32,
[[12, 5], [12, 5, 2, 2]], [[12, 5, 2, 2]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
expected_engines = {
"TRTEngineOp_000": [
"add0_1", "add0_2", "add0_3", "c0_1", "c0_2", "c0_3", "abs0_1",
"abs0_2", "expand0_0", "expand0_1", "axis"
],
"TRTEngineOp_001": [
"add1_1", "add1_2", "add1_3", "c1_1", "c1_2", "c1_3", "abs1_1",
"abs1_2", "reciprocal1"
]
}
if not run_params.dynamic_shape:
# The two ops can't be in the same cluster as the ops in TRTEngineOp_000
# due to trt_incompatible_op. They can't be in the same cluster as the
# ops in TRTEngineOP_1 because their batch size belongs to a different
# equivalent class.
expected_engines["TRTEngineOp_002"] = ["add", "reciprocal0"]
else:
# In dynamic shape mode the batch size of the ops can differ,
# therefore the final ops will be merged to TRTEngineOP_1.
expected_engines["TRTEngineOp_001"] += ["add", "reciprocal0"]
return expected_engines
if __name__ == "__main__":
test.main()
@@ -0,0 +1,126 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Basic tests for TF-TensorRT integration."""
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class ReshapeTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
outputs = []
# Here we test two types of reshapes, one changes the batch dimension and
# the other does not. Note that we're not able to test reshaping to
# scalar, since TRT requires input tensor to be of rank at least 2, so a
# reshape with scalar input will be filtered out of the segment before
# conversion.
#
# These reshapes happen at batch dimension, thus conversion should fail.
orig_shape = constant_op.constant([-1, 24, 24, 2], name="original_shape")
for shape in [[2, 50, 24, 24, 2], [-1, 50, 24, 24, 2], [2, 50, -1, 24, 2]]:
incompatible_reshape = array_ops.reshape(inp, shape)
reshape_back = array_ops.reshape(incompatible_reshape, orig_shape)
outputs.append(self.trt_incompatible_op(reshape_back))
# Add another block with many reshapes that don't change the batch
# dimension.
compatible_reshape = array_ops.reshape(
inp, [-1, 24 * 24, 2], name="reshape-0")
compatible_reshape = array_ops.reshape(
compatible_reshape, [100, 24, -1], name="reshape-1")
compatible_reshape = array_ops.reshape(
compatible_reshape, [100, 24 * 2, 24], name="reshape-2")
compatible_reshape = array_ops.reshape(
compatible_reshape, [-1, 24, 24 * 2], name="reshape-3")
compatible_reshape = array_ops.reshape(
compatible_reshape, [-1, 6, 4, 24, 2], name="reshape-4")
compatible_reshape = array_ops.reshape(
compatible_reshape, [-1, 6, 4, 6, 4, 2, 1], name="reshape-5")
compatible_reshape = array_ops.reshape(
compatible_reshape, [-1, 24, 24, 2], name="reshape-6")
outputs.append(self.trt_incompatible_op(compatible_reshape))
return math_ops.add_n(outputs, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
[[100, 24, 24, 2]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": ["reshape-%d" % i for i in range(7)] +
["reshape-%d/shape" % i for i in range(7)]
}
def ShouldRunTest(self, run_params):
"""Whether to run the test."""
return (not trt_test.IsQuantizationMode(run_params.precision_mode) and
not run_params.dynamic_engine), "test static engine and non-INT8"
class TransposeTest(trt_test.TfTrtIntegrationTestBase):
def GraphFn(self, inp):
# Add a block with compatible transposes.
compatible_transpose = array_ops.transpose(
inp, [0, 3, 1, 2], name="transpose-1")
compatible_transpose = array_ops.transpose(
compatible_transpose, [0, 2, 3, 1], name="transposeback")
return array_ops.identity(compatible_transpose, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
[[100, 24, 24, 2]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": [
"transpose-1", "transpose-1/perm", "transposeback",
"transposeback/perm"
]
}
def ShouldRunTest(self, run_params):
"""Whether to run the test."""
return (not trt_test.IsQuantizationMode(run_params.precision_mode) and
not run_params.dynamic_engine), "test static engine and non-INT8"
class IncompatibleTransposeTest(TransposeTest):
def GraphFn(self, inp):
# Add a block with incompatible transposes.
incompatible_transpose = array_ops.transpose(
inp, [2, 1, 0, 3], name="transpose-2")
excluded_transpose = array_ops.transpose(
incompatible_transpose, [0, 2, 3, 1], name="transpose-3")
return array_ops.identity(excluded_transpose, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
[[24, 100, 2, 24]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return []
if __name__ == "__main__":
test.main()
@@ -0,0 +1,282 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class ShapeOutputTest(trt_test.TfTrtIntegrationTestBase):
"""Test shape value output with TF-TRT."""
def setUp(self):
super().setUp()
self.DisableNonTrtOptimizers()
def GraphFn(self, x):
# The first engine returns the shape of q, which equals the shape of x. The
# values of x are actually not needed for the TRT engine. This way x is
# neither a shape tensor nor an execution tensor, we still need to set its
# shape (binding dimensions). We confirm with this test that the binding
# dimensions of x are correctly set before we execute the engine.
q = 2 * x + 1
q = array_ops.shape(q)
q = math_ops.cast(q, dtypes.float32)
q = self.trt_incompatible_op(q)
q = q * 2 + q * q
return array_ops.identity(q, name="output_0")
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[2, 2, 5, 3]], [[4]],
extra_inputs=[[[8, 2, 5, 3]]],
extra_outputs=[[[4]]],
input_mask=[[False, True, True, True]],
output_mask=[[True]])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
if run_params.dynamic_shape:
return ["TRTEngineOp_000", "TRTEngineOp_001"]
else:
# Second segment not converted in implicit batch mode, because its
# tensors have only one dimensions
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
# We cannot calibrate without building the engine, we turn of INT8 test.
return (run_params.dynamic_shape and
run_params.precision_mode != "INT8", "no calibration dynamic shape")
class ShapeOutputWithSingleInputProfile(ShapeOutputTest):
"""Same as the previous test, but with a single input profile."""
def setUp(self):
super().setUp()
self.DisableNonTrtOptimizers()
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[2, 2, 5, 3]], [[4]],
extra_inputs=[],
extra_outputs=[],
input_mask=[[False, True, True, True]],
output_mask=[[True]])
class ShapeOutputWithSingleInputAndReshape(trt_test.TfTrtIntegrationTestBase):
"""Similar to the previous test, but the ShapeOp output is reshaped to 2D.
This makes the output tensor not compatible with shape tensor.
"""
def setUp(self):
super().setUp()
self.DisableNonTrtOptimizers()
def GraphFn(self, x):
q = 2 * x + 1
q = array_ops.shape(q)
q = gen_array_ops.reshape(q, [2, 2])
q = math_ops.cast(q, dtypes.float32)
q = self.trt_incompatible_op(q)
q = q * 2 + q * q
return array_ops.identity(q, name="output_0")
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[2, 2, 5, 3]], [[2, 2]],
extra_inputs=[],
extra_outputs=[],
input_mask=[[False, True, True, True]],
output_mask=[[True, True]])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
return ["TRTEngineOp_000", "TRTEngineOp_001"]
class PrunedInputTest(trt_test.TfTrtIntegrationTestBase):
"""In TRT 7, an input tensor can be pruned if it is not used by the network.
This happens if only its shape is used, but the shape is already defined by
the optimization profile by setting min=max. (nvbugs/3153064)
After pruning, the TRT network has no input bindings.
"""
def setUp(self):
super().setUp()
self.DisableNonTrtOptimizers()
def GraphFn(self, x):
q = array_ops.shape(x)
q = q * 2 + q * q
return array_ops.identity(q, name="output_0")
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[1, 2, 5, 3]], [[4]],
extra_inputs=[],
extra_outputs=[],
input_mask=[[False, True, True, True]],
output_mask=[[True]])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
# Shape op is only converted in dynamic shape mode.
return (run_params.dynamic_shape and
run_params.is_v2, "test v2 dynamic shape")
class PrunedInputTest2(trt_test.TfTrtIntegrationTestBase):
"""Two inputs, one of the is pruned."""
def setUp(self):
super().setUp()
self.DisableNonTrtOptimizers()
def GraphFn(self, x, y):
q = array_ops.shape(x)
z = y * y + y
z = gen_array_ops.reshape(z, q)
out_0 = array_ops.identity(q, name="output_0")
out_1 = array_ops.identity(z, name="output_1")
return (out_0, out_1)
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[1, 2, 5, 3], [2, 15]], [[4], [1, 2, 5, 3]],
extra_inputs=[],
extra_outputs=[],
input_mask=[[False, True, True, True], [False, True]],
output_mask=[[True], [False, True, True, True]])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
# Shape op is only converted in dynamic shape mode.
return (run_params.dynamic_shape and
run_params.is_v2, "test v2 dynamic shape")
class ShapeValueMaskTest(trt_test.TfTrtIntegrationTestBase):
"""Confirm that 0D and 1D non int tensors are not treated as shape tensors."""
def setUp(self):
super().setUp()
# This is to test whether shape value mask is correctly set in case engine
# construction has failed.
os.environ["TF_TRT_ABORT_CUDA_ENGINE_BUILD"] = "True"
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
def tearDown(self):
super().tearDown()
os.environ["TF_TRT_ABORT_CUDA_ENGINE_BUILD"] = "False"
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "False"
def GraphFn(self, x, y):
q = 2 * x + y
return array_ops.identity(q, name="output_0")
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float16, [[3], []], [[3]],
extra_inputs=[],
extra_outputs=[],
input_mask=[[True], []],
output_mask=[[True]])
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
if run_params.dynamic_shape:
return ["TRTEngineOp_000"]
else:
return []
def ShouldRunTest(self, run_params):
# We cannot calibrate without building the engine, we turn of INT8 test.
return (run_params.dynamic_shape and
run_params.precision_mode != "INT8", "no calibration dynamic shape")
class InputProfile(trt_test.TfTrtIntegrationTestBase):
"""The shape profiles has to fit values of shape tensors, but for regular
tensors the values do not matter. Here we test shape profile management with
an INT32 input tensor that is not a shape tensor. The extra inputs with
dim=10 would trigger an error if we mistakenly treat it as shape tensors.
"""
def setUp(self):
super().setUp()
self.DisableNonTrtOptimizers()
def GraphFn(self, x):
z = x * x + x + 1
z = array_ops.identity(z, name="output_0")
return z
def GetParams(self):
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.int32,
[[4]],
[[4]],
extra_inputs=[[[5]], [[10]]],
extra_outputs=[[[5]], [[10]]],
input_mask=[[False]],
output_mask=[[False]],
)
def ExpectedEnginesToBuild(self, run_params):
"""Returns the expected engines to build."""
return ["TRTEngineOp_000"]
def ShouldRunTest(self, run_params):
# Shape op is only converted in dynamic shape mode.
return (
run_params.dynamic_shape
and run_params.is_v2
and not trt_test.IsQuantizationMode(run_params.precision_mode),
"Test v2 dynamic_shapes without INT8",
)
if __name__ == "__main__":
test.main()
@@ -0,0 +1,49 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities to test TF-TensorRT integration."""
import os
from contextlib import contextmanager
@contextmanager
def experimental_feature_scope(feature_name):
"""Creates a context manager to enable the given experimental feature.
This helper function creates a context manager setting up an experimental
feature temporarily.
Example:
```python
with self._experimental_feature_scope("feature_1"):
do_smthg()
```
Args:
feature_name: Name of the feature being tested for activation.
"""
env_varname = "TF_TRT_EXPERIMENTAL_FEATURES"
env_value_bckp = os.environ.get(env_varname, default="")
exp_features = env_value_bckp.split(",")
os.environ[env_varname] = ",".join(list(set(exp_features + [feature_name])))
try:
yield
finally:
os.environ[env_varname] = env_value_bckp
@@ -0,0 +1,61 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""Saves a TensorFlow model containing ReadVariableOp nodes.
The saved model is loaded and executed by tests to check that TF-TRT can
successfully convert and execute models with variables without freezing.
"""
import numpy as np
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import save
class MyModel(module.Module):
"""Simple model with two variables."""
def __init__(self):
self.var1 = variables.Variable(
np.array([[[13.]]], dtype=np.float32), name="var1")
self.var2 = variables.Variable(
np.array([[[37.]]], dtype=np.float32), name="var2")
@def_function.function
def __call__(self, input1, input2):
mul1 = input1 * self.var1
mul2 = input2 * self.var2
add = mul1 + mul2
sub = add - 45.
return array_ops.identity(sub, name="output")
def GenerateModelWithReadVariableOp(tf_saved_model_dir):
"""Generate a model with ReadVariableOp nodes."""
my_model = MyModel()
cfunc = my_model.__call__.get_concrete_function(
tensor_spec.TensorSpec([None, 1, 1], dtypes.float32),
tensor_spec.TensorSpec([None, 1, 1], dtypes.float32))
# pylint: disable=not-callable
save(my_model, tf_saved_model_dir, signatures=cfunc)
if __name__ == "__main__":
GenerateModelWithReadVariableOp(
tf_saved_model_dir="tf_readvariableop_saved_model")
@@ -0,0 +1,90 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""Saves a TensorFlow model containing VariableV2 nodes.
The saved model is loaded and executed by tests to check that TF-TRT can
successfully convert and execute models with VariableV2 without freezing.
"""
from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import builder
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model import utils
ops.disable_eager_execution()
variable_scope.disable_resource_variables()
def GenerateModelWithVariableV2(tf_saved_model_dir):
"""Generate a model with a VariableV2 node using TFv1 API."""
def SimpleModel():
"""Define model with a TF graph."""
def GraphFn():
input1 = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input1")
input2 = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input2")
var1 = variable_scope.get_variable(
"var1",
shape=[1, 1, 1],
initializer=init_ops.constant_initializer([[[13.]]]),
dtype=dtypes.float32)
var2 = variable_scope.get_variable(
"var2",
shape=[1, 1, 1],
initializer=init_ops.constant_initializer([[[37.]]]),
dtype=dtypes.float32)
mul1 = input1 * var1
mul2 = input2 * var2
add = mul1 + mul2
sub = add - 45.
out = array_ops.identity(sub, name="output")
return g, input1, input2, out
g = ops.Graph()
with g.as_default():
return GraphFn()
g, input1, input2, out = SimpleModel()
signature_def = signature_def_utils.build_signature_def(
inputs={
"input1": utils.build_tensor_info(input1),
"input2": utils.build_tensor_info(input2)
},
outputs={"output": utils.build_tensor_info(out)},
method_name=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY)
saved_model_builder = builder.SavedModelBuilder(tf_saved_model_dir)
with session.Session(graph=g) as sess:
variables.global_variables_initializer().run()
saved_model_builder.add_meta_graph_and_variables(
sess, [tag_constants.SERVING],
signature_def_map={
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
})
saved_model_builder.save()
if __name__ == "__main__":
GenerateModelWithVariableV2(tf_saved_model_dir="tf_variablev2_saved_model")
@@ -0,0 +1,141 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""Saves a SavedModel after TensorRT conversion.
The saved model is loaded and executed by tests to ensure backward
compatibility across TF versions.
The script may not work in TF1.x.
Instructions on how to use this script:
- Execute the script as follows:
python gen_tftrt_model
- Rename tftrt_saved_model to what makes sense for your test.
- Delete directory tf_saved_model unless you want to use it.
"""
from tensorflow.python import Session
from tensorflow.python.compiler.tensorrt import trt_convert
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import builder
from tensorflow.python.saved_model import save
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model import utils
from tensorflow.python.trackable import autotrackable
def GetGraph(input1, input2, var):
"""Define graph."""
add = input1 + var
mul = input1 * add
add = mul + add
add = add + input2
out = array_ops.identity(add, name="output")
return out
def GenerateModelV2(tf_saved_model_dir, tftrt_saved_model_dir):
"""Generate and convert a model using TFv2 API."""
class SimpleModel(autotrackable.AutoTrackable):
"""Define model with a TF function."""
def __init__(self):
self.v = None
@def_function.function(input_signature=[
tensor_spec.TensorSpec(shape=[None, 1, 1], dtype=dtypes.float32),
tensor_spec.TensorSpec(shape=[None, 1, 1], dtype=dtypes.float32)
])
def run(self, input1, input2):
if self.v is None:
self.v = variables.Variable([[[1.0]]], dtype=dtypes.float32)
return GetGraph(input1, input2, self.v)
root = SimpleModel()
# Saved TF model
# pylint: disable=not-callable
save(
root,
tf_saved_model_dir,
{signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: root.run})
# Convert TF model to TensorRT
converter = trt_convert.TrtGraphConverterV2(
input_saved_model_dir=tf_saved_model_dir)
converter.convert()
try:
line_length = max(160, os.get_terminal_size().columns)
except OSError:
line_length = 160
converter.summary(line_length=line_length, detailed=True)
converter.save(tftrt_saved_model_dir)
def GenerateModelV1(tf_saved_model_dir, tftrt_saved_model_dir):
"""Generate and convert a model using TFv1 API."""
def SimpleModel():
"""Define model with a TF graph."""
def GraphFn():
input1 = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input1")
input2 = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input2")
var = variables.Variable([[[1.0]]], dtype=dtypes.float32, name="v1")
out = GetGraph(input1, input2, var)
return g, var, input1, input2, out
g = ops.Graph()
with g.as_default():
return GraphFn()
g, var, input1, input2, out = SimpleModel()
signature_def = signature_def_utils.build_signature_def(
inputs={
"input1": utils.build_tensor_info(input1),
"input2": utils.build_tensor_info(input2)
},
outputs={"output": utils.build_tensor_info(out)},
method_name=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY)
saved_model_builder = builder.SavedModelBuilder(tf_saved_model_dir)
with Session(graph=g) as sess:
sess.run(var.initializer)
saved_model_builder.add_meta_graph_and_variables(
sess, [tag_constants.SERVING],
signature_def_map={
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
})
saved_model_builder.save()
# Convert TF model to TensorRT
converter = trt_convert.TrtGraphConverter(
input_saved_model_dir=tf_saved_model_dir, is_dynamic_op=True)
converter.convert()
converter.save(tftrt_saved_model_dir)
if __name__ == "__main__":
GenerateModelV2(
tf_saved_model_dir="tf_saved_model",
tftrt_saved_model_dir="tftrt_saved_model")
@@ -0,0 +1,3 @@
model_checkpoint_path: "model.ckpt-46900"
all_model_checkpoint_paths: "model.ckpt-0"
all_model_checkpoint_paths: "model.ckpt-46900"
@@ -0,0 +1,348 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""TF function conversion."""
import itertools
import os
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.protobuf import saved_model_pb2
from tensorflow.python.compiler.tensorrt import trt_convert
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.compiler.tensorrt.test.tf_trt_integration_test_base import GraphState
from tensorflow.python.compiler.tensorrt.test.tf_trt_integration_test_base import IsQuantizationWithCalibration
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import constants
from tensorflow.python.saved_model import loader_impl
from tensorflow.python.util import compat
class TfFunctionTest(trt_test.TfTrtIntegrationTestBase):
def __init__(self, methodName): # pylint: disable=invalid-name
super(TfFunctionTest, self).__init__(methodName)
self._profile_strategy = "Range"
self._trt_engine_op_count_offset = 0
self._test_conversion_params = {
"_tftrt_convert_function": True,
"_tftrt_trt_logger_name": "DefaultLogger",
"_tftrt_max_batch_size": 10,
"_tftrt_max_workspace_size_bytes":
(trt_convert.DEFAULT_TRT_MAX_WORKSPACE_SIZE_BYTES),
"_tftrt_precision_mode": "FP16",
"_tftrt_minimum_segment_size": 2,
"_tftrt_is_dyn_op": True,
"_tftrt_max_cached_engines": 1,
"_tftrt_use_calibration": False,
"_tftrt_use_implicit_batch": True,
"_tftrt_profile_strategy": self._profile_strategy,
"_tftrt_allow_build_at_runtime": False
}
self._is_v2 = False
def ShouldRunTest(self, run_params):
should_run, reason_for_skipping = (
trt_test.TfTrtIntegrationTestBase.ShouldRunTest(self, run_params))
if not should_run:
return should_run, reason_for_skipping
else:
# TODO(kyungtaek): Calibration currently does not run for nodes
# nested within functions. If this gets fixed, this method should not
# override the parent method.
return (not IsQuantizationWithCalibration(run_params),
"calibration is not supported for tf.functions")
def _copy_test_attr_to_func_def(self, func_def, param_name, attr_value_type):
test_value = self._test_conversion_params[param_name]
if attr_value_type == "s":
byte_value = compat.as_bytes(test_value)
func_def.attr[param_name].CopyFrom(attr_value_pb2.AttrValue(s=byte_value))
elif attr_value_type == "b":
func_def.attr[param_name].CopyFrom(attr_value_pb2.AttrValue(b=test_value))
elif attr_value_type == "i":
func_def.attr[param_name].CopyFrom(attr_value_pb2.AttrValue(i=test_value))
else:
logging.info("Attr_value type %s is not supported", attr_value_type)
def _ChainAllNodes(self, graph_def):
return itertools.chain(
graph_def.node,
itertools.chain(
*[func.node_def for func in graph_def.library.function]))
def _VerifyTestAttrs(self, function_protos):
if self._test_conversion_params["_tftrt_convert_function"]:
for func_def in function_protos:
if not func_def.signature.name.startswith("TRTEngine"):
for key, value in self._test_conversion_params.items():
self.assertIn(key, func_def.attr,
"key %s not found in func_def.attr" % key)
if isinstance(value, str):
self.assertEqual(func_def.attr[key].s, compat.as_bytes(value))
elif isinstance(value, bool):
self.assertEqual(func_def.attr[key].b, value)
elif isinstance(value, int):
self.assertEqual(func_def.attr[key].i, value)
@def_function.function(input_signature=[
tensor_spec.TensorSpec(shape=[None, 32, 32, 2], dtype=dtypes.float32)
])
def _conv_and_pool_0(self, inp):
dtype = inp.dtype
conv_filter = constant_op.constant([[[[1., 0.5, 4.], [1., 0.5, 1.]]]],
name="weights",
dtype=dtype)
conv = nn.conv2d(
input=inp,
filter=conv_filter,
strides=[1, 2, 2, 1],
padding="SAME",
name="conv")
bias = constant_op.constant([4., 1.5, 2.], name="bias", dtype=dtype)
added = nn.bias_add(conv, bias, name="bias_add")
relu = nn.relu(added, "relu")
identity = array_ops.identity(relu, "identity")
pool = nn_ops.max_pool(
identity, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
return array_ops.squeeze(pool)
def GraphFn(self, x):
x = self._conv_and_pool_0(x)
return array_ops.identity(x, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[10, 32, 32, 2]],
[[10, 8, 8, 3]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {
"TRTEngineOp_000": [
"weights", "conv", "bias", "bias_add", "relu", "identity",
"max_pool"
]
}
def _copy_test_attributes_to_func_def(self, func_def):
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_convert_function",
attr_value_type="b")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_trt_logger_name",
attr_value_type="s")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_max_batch_size",
attr_value_type="i")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_max_workspace_size_bytes",
attr_value_type="i")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_precision_mode",
attr_value_type="s")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_minimum_segment_size",
attr_value_type="i")
self._copy_test_attr_to_func_def(
func_def=func_def, param_name="_tftrt_is_dyn_op", attr_value_type="b")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_max_cached_engines",
attr_value_type="i")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_use_calibration",
attr_value_type="b")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_use_implicit_batch",
attr_value_type="b")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_profile_strategy",
attr_value_type="s")
self._copy_test_attr_to_func_def(
func_def=func_def,
param_name="_tftrt_allow_build_at_runtime",
attr_value_type="b")
def _MakeSavedModelV1(self, run_params):
"""Write the saved model as an input for testing.
In addition to creating a SavedModel like its parent method, this method
replaces this SavedModel by adding TF-TRT conversion parameters as function
attributes to each function in the SavedModel.
Args:
run_params: The current test run parameters.
Returns:
The directory of the saved model.
"""
saved_model_dir = trt_test.TfTrtIntegrationTestBase._MakeSavedModelV1(
self, run_params)
saved_model_proto = loader_impl.parse_saved_model(saved_model_dir)
new_saved_model = saved_model_pb2.SavedModel()
new_saved_model.CopyFrom(saved_model_proto)
new_meta_graph_def = new_saved_model.meta_graphs[0]
for func_def in new_meta_graph_def.graph_def.library.function:
# Disable function inlining.
func_def.attr["_noinline"].CopyFrom(attr_value_pb2.AttrValue(b=True))
self._copy_test_attributes_to_func_def(func_def)
old_saved_model_file = os.path.join(saved_model_dir,
constants.SAVED_MODEL_FILENAME_PB)
if os.path.exists(old_saved_model_file):
os.remove(old_saved_model_file)
path = os.path.join(
compat.as_bytes(saved_model_dir),
compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
file_io.write_string_to_file(
path, new_saved_model.SerializeToString(deterministic=True))
return saved_model_dir
def _MakeSavedModelV2(self, run_params):
"""Write the saved model as an input for testing.
In addition to creating a SavedModel like its parent method, this method
replaces this SavedModel by adding TF-TRT conversion parameters as function
attributes to each function in the SavedModel.
Args:
run_params: The current test run parameters.
Returns:
The directory of the saved model.
"""
saved_model_dir = trt_test.TfTrtIntegrationTestBase._MakeSavedModelV2(
self, run_params)
saved_model_proto = loader_impl.parse_saved_model(saved_model_dir)
new_saved_model = saved_model_pb2.SavedModel()
new_saved_model.CopyFrom(saved_model_proto)
new_meta_graph_def = new_saved_model.meta_graphs[0]
prefix_len = len("__inference_")
for func_def in new_meta_graph_def.graph_def.library.function:
logging.info("_MakeSavedModelV2, func_def name: %s",
func_def.signature.name)
func_name_without_prefix = func_def.signature.name[prefix_len:]
if func_name_without_prefix.startswith(
("_conv_and_pool_0")):
func_def.attr["_noinline"].CopyFrom(attr_value_pb2.AttrValue(b=True))
self._copy_test_attributes_to_func_def(func_def)
old_saved_model_file = os.path.join(saved_model_dir,
constants.SAVED_MODEL_FILENAME_PB)
if os.path.exists(old_saved_model_file):
os.remove(old_saved_model_file)
path = os.path.join(
compat.as_bytes(saved_model_dir),
compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
file_io.write_string_to_file(
path, new_saved_model.SerializeToString(deterministic=True))
return saved_model_dir
def _VerifyGraphDefV1(self, run_params, original_gdef, gdef_to_verify,
graph_state):
expected_engines = self.ExpectedEnginesToBuild(run_params)
num_engines = 0
functions = [f.signature.name for f in gdef_to_verify.library.function]
all_nodes = list(self._ChainAllNodes(gdef_to_verify))
all_nodes.sort(key=lambda x: x.name)
for node in all_nodes:
if node.op == "TRTEngineOp":
logging.info("Found TRTEngineOp: " + node.name)
num_engines += 1
segment_funcdef_name = node.attr["segment_func"].func.name
function_name = node.name + "_native_segment"
is_dynamic_engine = not node.attr["static_engine"].b
self.assertNotEmpty(segment_funcdef_name, node.name)
self.assertIn(function_name, functions)
if (not IsQuantizationWithCalibration(run_params) and
not is_dynamic_engine):
self.assertTrue(len(node.attr["serialized_segment"].s), node.name)
self.assertIn(
self._RemoveGraphSequenceNumber(node.name), expected_engines)
self.assertEqual(
self._ToBytes(run_params.precision_mode),
node.attr["precision_mode"].s, node.name)
self.assertEqual(run_params.dynamic_engine, is_dynamic_engine,
node.name)
self.assertEqual(node.attr["use_calibration"].b,
run_params.use_calibration, node.name)
has_calibration_data = len(node.attr["calibration_data"].s)
if (IsQuantizationWithCalibration(run_params) and
graph_state == GraphState.INFERENCE):
self.assertTrue(has_calibration_data, node.name)
else:
self.assertFalse(has_calibration_data, node.name)
if graph_state == GraphState.ORIGINAL:
self.assertEqual(0, num_engines)
self._VerifyTestAttrs(function_protos=gdef_to_verify.library.function)
else:
self.assertEqual(num_engines, len(expected_engines))
expected_connections = self.ExpectedConnections(run_params)
if expected_connections:
self._VerifyConnections(expected_engines, expected_connections,
original_gdef, gdef_to_verify)
self._VerifyMaxBatchSizeAnnotations(
expected_engines=expected_engines,
original_gdef=original_gdef,
converted_gdef=gdef_to_verify,
expected_max_batch_sizes=self.ExpectedMaxBatchSizes(run_params),
default_max_batch_size=self.GetMaxBatchSize(run_params))
self._VerifyTestAttrs(function_protos=gdef_to_verify.library.function)
def _ShouldConverterBuild(self, run_params):
return (run_params.is_v2 and not run_params.convert_online and
run_params.dynamic_engine)
def RunTest(self, run_params):
self._test_conversion_params["_tftrt_precision_mode"] = (
run_params.precision_mode)
self._test_conversion_params["_tftrt_use_calibration"] = (
run_params.use_calibration)
self._test_conversion_params["_tftrt_is_dyn_op"] = (
run_params.dynamic_engine)
# When running with V1, using dynamic_engine and
# allow_build_at_runtime==False at the same time do not work.
if run_params.is_v2:
self._test_conversion_params["_tftrt_allow_build_at_runtime"] = True
self._is_v2 = True
else:
self._test_conversion_params["_tftrt_allow_build_at_runtime"] = (
run_params.convert_online or run_params.dynamic_engine)
self._test_conversion_params["_tftrt_use_implicit_batch"] = \
not run_params.dynamic_shape
self.DisableNonTrtOptimizers()
trt_test.TfTrtIntegrationTestBase.RunTest(self, run_params)
if __name__ == "__main__":
test.main()
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,72 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import constant_op
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
class TopKTest(trt_test.TfTrtIntegrationTestBase):
"""Testing Top-K in TF-TRT conversion."""
def GraphFn(self, x):
k = 5
k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
values = array_ops.identity(values, name="output_0")
indices = array_ops.identity(indices, name="output_1")
return values, indices
def GetParams(self):
k = 5
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 100]],
[[100, k], [100, k]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {"TRTEngineOp_000": ["Const", "TopK"]}
class TopKOutputTypeTest(trt_test.TfTrtIntegrationTestBase):
"""Testing that output type of engine using Top-K is set correctly."""
def GraphFn(self, x):
k = 5
k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
# Reshape will act as a layer between the TopK output and the engine
# output, requiring the output tensor of reshape to be set explicitly to
# int32.
indices = array_ops.reshape(indices, [100, 1, 5], name="Reshape")
values = array_ops.identity(values, name="output_0")
indices = array_ops.identity(indices, name="output_1")
return values, indices
def GetParams(self):
k = 5
return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 100]],
[[100, k], [100, 1, k]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return {"TRTEngineOp_000": ["Const", "TopK", "Reshape", "Reshape/shape"]}
if __name__ == "__main__":
test.main()
@@ -0,0 +1,93 @@
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""This script test input and output shapes and dtype of the TRTEngineOp."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
from tensorflow.python.saved_model import load
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
class TRTEngineOpInputOutputShapeTest(trt_test.TfTrtIntegrationTestBase):
"""Testing the output shape of a TRTEngine."""
def GraphFn(self, inp):
b = array_ops.squeeze(inp, axis=[2])
c = nn.relu(b)
d1 = c + c
d2 = math_ops.reduce_sum(d1)
d1 = array_ops.identity(d1, name="output_0")
d2 = array_ops.identity(d2, name="output_1")
return d1, d2
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[1, 2, 1, 4]],
[[1, 2, 4], []])
def _GetInferGraph(self, *args, **kwargs):
trt_saved_model_dir = super(TRTEngineOpInputOutputShapeTest,
self)._GetInferGraph(*args, **kwargs)
def get_func_from_saved_model(saved_model_dir):
try: # Necessary for `bazel run ...`
saved_model_load_fn = load.load
except AttributeError: # All the other cases
saved_model_load_fn = load
saved_model_loaded = saved_model_load_fn(
saved_model_dir, tags=[tag_constants.SERVING])
graph_func = saved_model_loaded.signatures[
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
return graph_func, saved_model_loaded
func, _ = get_func_from_saved_model(trt_saved_model_dir)
input_shape = func.inputs[0].shape
if isinstance(input_shape, tensor_shape.TensorShape):
input_shape = input_shape.as_list()
output_shapes = [
out_shape.shape.as_list() if isinstance(
out_shape.shape, tensor_shape.TensorShape) else out_shape.shape
for out_shape in func.outputs
]
self.assertEqual(func.inputs[0].dtype, dtypes.float32)
self.assertEqual(func.outputs[0].dtype, dtypes.float32)
self.assertEqual(func.outputs[1].dtype, dtypes.float32)
self.assertEqual(input_shape, [None, 2, 1, 4])
self.assertEqual(output_shapes[0], [None, 2, 4])
self.assertEqual(output_shapes[1], [])
return trt_saved_model_dir
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,131 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
from unittest import SkipTest # pylint: disable=g-importing-member
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class TrtModeTestBase(trt_test.TfTrtIntegrationTestBase):
"""Test squeeze on batch dim and some unary operations in TF-TRT."""
def GraphFn(self, x1):
q = math_ops.abs(x1)
q = q + 1.0
q = q * 3.0
q = array_ops.squeeze(q, 0)
q = math_ops.abs(q)
q = q + 5.0
return array_ops.identity(q, name="output_0")
def ShouldRunTest(self, run_params):
# Squeeze op produces dynamic shaped values. Therefore, we don't run the
# test with static engine to avoid native segment execution.
return (run_params.dynamic_engine and run_params.is_v2 and
not run_params.use_calibration, "test v2 dynamic engine and "
"non-calibration")
def GetParams(self):
"""The input has 1 as a first dimension, which is removed by the squeeze.
op in the graph.
In explicit batch mode, TensorRT can convert the whole graph. In this mode
it is possible to manipulate the batch dimension using the squeeze op.
In implicit batch mode TensorRT cannot convert the whole graph. We are not
allowed to manipulate (squeeze) the first dimension in implicit batch mode.
Therefore the graph will be converted using multiple segments.
"""
return self.BuildParams(self.GraphFn, dtypes.float32, [[1, 12, 5]],
[[12, 5]])
def GetMaxBatchSize(self, run_params):
if run_params.dynamic_engine:
return None
# The first dimension of the input is squeezed and the batch size for the
# rest OPs is 12.
return 12
@classmethod
def setUpClass(cls):
if cls is TrtModeTestBase:
raise SkipTest("TrtModeTestBase defines base class for other test.")
super(TrtModeTestBase, cls).setUpClass()
def ExpectedEnginesToBuild(self, run_params):
"""Check that the expected engine is built.
Args:
run_params: the run parameters.
Returns:
the expected engines to build.
The squeeze op is not converted by TensorRT in implicit batch mode.
Because of this we have two TRTEngineOp in the graphs: one for the
subgraph before 'squeeze(q,0)', and another one for the rest of the ops
after the 'squeeze(q,0)'.
In explicit batch mode the whole graph is converted using a single engine.
"""
if run_params.dynamic_shape:
return ["TRTEngineOp_000"]
else:
return ["TRTEngineOp_000", "TRTEngineOp_001"]
class StaticInputTest(TrtModeTestBase):
def GetParams(self):
"""We specify input/output masks with static (known) shapes."""
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[1, 12, 5]], [[12, 5]],
input_mask=[[True, True, True]],
output_mask=[[True, True]],
extra_inputs=[],
extra_outputs=[])
class DynamicInputTest(TrtModeTestBase):
"""Test with dynamic input shapes.
The difference to the previous test is that we use input and output masks to
change the input and output shapes to unknown shapes.
"""
def GetParams(self):
"""We specify input/output mask with dynamic (unknown) shapes.
In dynamic shape mode, single engine with three optimization profiles can
handle the three different input shapes.
"""
return self.BuildParamsWithMask(
self.GraphFn,
dtypes.float32, [[1, 12, 5]], [[12, 5]],
extra_inputs=[[[1, 2, 3]], [[1, 4, 6]]],
extra_outputs=[[[2, 3]], [[4, 6]]],
input_mask=[[False, False, False]],
output_mask=[[False, False]])
if __name__ == "__main__":
test.main()
@@ -0,0 +1,106 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class UnaryTest(trt_test.TfTrtIntegrationTestBase):
"""Test for unary operations in TF-TRT."""
def GraphFn(self, x1, x2):
x = x1
q = math_ops.abs(x)
q = q + 1.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = array_ops.squeeze(q, axis=-2)
q = math_ops.abs(q)
q = q + 2.2
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = array_ops.squeeze(q, axis=3)
q = math_ops.abs(q)
q = q + 3.0
a = gen_math_ops.reciprocal(q)
# this chain of operations has a batch size of 5, which is different from
# the batch size for the other operations.
x = constant_op.constant(np.random.randn(5, 8, 12), dtype=x.dtype)
q = math_ops.abs(x)
q = q + 2.0
q = gen_math_ops.exp(q)
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 2.1
q = gen_math_ops.sqrt(q)
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 4.0
b = gen_math_ops.reciprocal(q)
# TODO(jie): this one will break, broadcasting on batch.
x = x2
q = math_ops.abs(x)
q = q + 5.0
q = gen_math_ops.exp(q)
q = array_ops.squeeze(q, axis=[-1, -2, 3])
q = gen_math_ops.log(q)
q = math_ops.abs(q)
q = q + 5.1
q = gen_array_ops.reshape(q, [12, 5, 1, 1, 8, 1, 12])
q = array_ops.squeeze(q, axis=[5, 2, 3])
q = gen_math_ops.sqrt(q)
q = math_ops.abs(q)
q = q + 5.2
q = gen_math_ops.rsqrt(q)
q = math_ops.negative(q)
q = math_ops.abs(q)
q = q + 5.3
c = gen_math_ops.reciprocal(q)
q = a * b
q = q / c
return array_ops.squeeze(q, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32,
[[12, 5, 8, 1, 1, 12], [12, 5, 8, 1, 12, 1, 1]],
[[12, 5, 8, 12]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
if run_params.dynamic_shape:
# All the ops are converted into a single TRTEngineOp.
return ["TRTEngineOp_000"]
else:
# Final squeeze and div ops are not converted. The x1 and x2 branches
# are segmented into separate TRTEngineOp.
return ["TRTEngineOp_000", "TRTEngineOp_001"]
if __name__ == "__main__":
test.main()
@@ -0,0 +1,77 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import os
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_impl
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
class VGGBlockNCHWTest(trt_test.TfTrtIntegrationTestBase):
"""Single vgg layer in NCHW unit tests in TF-TRT."""
def GraphFn(self, x):
dtype = x.dtype
x, _, _ = nn_impl.fused_batch_norm(
x, [1.0, 1.0], [0.0, 0.0],
mean=[0.5, 0.5],
variance=[1.0, 1.0],
data_format="NCHW",
is_training=False)
e = constant_op.constant(
np.random.randn(1, 1, 2, 6), name="weights", dtype=dtype)
conv = nn.conv2d(
input=x,
filter=e,
data_format="NCHW",
strides=[1, 1, 2, 2],
padding="SAME",
name="conv")
b = constant_op.constant(np.random.randn(6), name="bias", dtype=dtype)
t = nn.bias_add(conv, b, data_format="NCHW", name="biasAdd")
relu = nn.relu(t, "relu")
idty = array_ops.identity(relu, "ID")
v = nn_ops.max_pool(
idty, [1, 1, 2, 2], [1, 1, 2, 2],
"VALID",
data_format="NCHW",
name="max_pool")
return array_ops.squeeze(v, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[5, 2, 8, 8]],
[[5, 6, 2, 2]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
# TODO(b/159459919): remove this routine to disallow native segment execution.
def setUp(self):
super().setUp()
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
if __name__ == "__main__":
test.main()
@@ -0,0 +1,68 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
import os
import numpy as np
from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_impl
from tensorflow.python.ops import nn_ops
from tensorflow.python.platform import test
class VGGBlockTest(trt_test.TfTrtIntegrationTestBase):
"""Single vgg layer test in TF-TRT conversion."""
def GraphFn(self, x):
dtype = x.dtype
x, _, _ = nn_impl.fused_batch_norm(
x, [1.0, 1.0], [0.0, 0.0],
mean=[0.5, 0.5],
variance=[1.0, 1.0],
is_training=False)
e = constant_op.constant(
np.random.randn(1, 1, 2, 6), name="weights", dtype=dtype)
conv = nn.conv2d(
input=x, filter=e, strides=[1, 2, 2, 1], padding="SAME", name="conv")
b = constant_op.constant(np.random.randn(6), name="bias", dtype=dtype)
t = nn.bias_add(conv, b, name="biasAdd")
relu = nn.relu(t, "relu")
idty = array_ops.identity(relu, "ID")
v = nn_ops.max_pool(
idty, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
return array_ops.squeeze(v, name="output_0")
def GetParams(self):
return self.BuildParams(self.GraphFn, dtypes.float32, [[5, 8, 8, 2]],
[[5, 2, 2, 6]])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_000"]
# TODO(b/159459919): remove this routine to disallow native segment execution.
def setUp(self):
super().setUp()
os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
if __name__ == "__main__":
test.main()