text: # 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 __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  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_0"]
  
    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)
  
    def ShouldRunTest(self, run_params):
      # TODO(b/162448349): Enable the test for TRT 7.1.3.
      if trt_test.IsTensorRTVersionGreaterEqual(7, 1, 3):
        return (False, "Skip test due to b/162448349")
      return super().ShouldRunTest(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_0"]
  
    def ShouldRunTest(self, run_params):
      # TODO(b/162448349): Enable the test for TRT 7.1.3.
      if trt_test.IsTensorRTVersionGreaterEqual(7, 1, 3):
        return (False, "Skip test due to b/162448349")
      return super().ShouldRunTest(run_params)
  
  
  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_0"]
  
    def ShouldRunTest(self, run_params):
      # TODO(b/162448349): Enable the test for TRT 7.1.3.
      if trt_test.IsTensorRTVersionGreaterEqual(7, 1, 3):
        return (False, "Skip test due to b/162448349")
      return super().ShouldRunTest(run_params)
  
  
  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_0"]
  
  
  if __name__ == "__main__":
    test.main()
  