# Copyright (c) 2021 PaddlePaddle 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.

from __future__ import print_function

import unittest
import numpy as np
from inference_pass_test import InferencePassTest
import paddle.fluid as fluid
import paddle.fluid.core as core
from paddle.fluid.core import PassVersionChecker
from paddle.fluid.core import AnalysisConfig


class TRTReduceMeanTest(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[-1, 3, -1, -1], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(
                data, dim=[2, -1], keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([3, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanTest.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, False, False)
        self.fetch_list = [out]
        self.dynamic_shape_params = TRTReduceMeanTest.DynamicShapeParam({
            'data': [1, 3, 16, 16]
        }, {'data': [3, 3, 56, 56]}, {'data': [3, 3, 56, 56]}, False)

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanAllNoBatchTest(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[-1, 3, -1, -1], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(data, keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([3, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanAllNoBatchTest.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, False, False)
        self.fetch_list = [out]
        self.dynamic_shape_params = TRTReduceMeanAllNoBatchTest.DynamicShapeParam(
            {
                'data': [1, 3, 16, 16]
            }, {'data': [3, 3, 56, 56]}, {'data': [3, 3, 56, 56]}, False)

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanTestFP16(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[-1, 3, -1, -1], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(
                data, dim=[2, -1], keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([3, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanTestFP16.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Half, False, False)
        self.fetch_list = [out]
        self.dynamic_shape_params = TRTReduceMeanTestFP16.DynamicShapeParam({
            'data': [1, 3, 16, 16]
        }, {'data': [3, 3, 56, 56]}, {'data': [3, 3, 56, 56]}, False)

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanAllTest(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[-1, 3, 56, 56], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(data, keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([3, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanAllTest.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, False, False)
        self.fetch_list = [out]
        self.dynamic_shape_params = TRTReduceMeanAllTest.DynamicShapeParam({
            'data': [1, 3, 56, 56]
        }, {'data': [3, 3, 56, 56]}, {'data': [3, 3, 56, 56]}, False)

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanTestStatic(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[3, 3, 56, 56], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(
                data, dim=[2, -1], keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([3, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanTestStatic.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, False, False)
        self.fetch_list = [out]

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanStaticAllTest(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[4, 3, 56, 56], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(data, keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([4, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanStaticAllTest.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, False, False)
        self.fetch_list = [out]

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanStaticFP16(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[4, 3, 56, 56], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(data, keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([4, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanStaticFP16.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Half, False, False)
        self.fetch_list = [out]

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


class TRTReduceMeanFP16Static(InferencePassTest):
    def setUp(self):
        with fluid.program_guard(self.main_program, self.startup_program):
            data = fluid.data(
                name="data", shape=[4, 3, 56, 56], dtype="float32")
            reduce_mean = fluid.layers.reduce_mean(data, keep_dim=True)
            out = fluid.layers.batch_norm(reduce_mean, is_test=True)

        self.feeds = {
            "data": np.random.random([4, 3, 56, 56]).astype("float32"),
        }
        self.enable_trt = True
        self.trt_parameters = TRTReduceMeanFP16Static.TensorRTParam(
            1 << 30, 32, 1, AnalysisConfig.Precision.Half, True, False)
        self.fetch_list = [out]

    def test_check_output(self):
        if core.is_compiled_with_cuda():
            use_gpu = True
            self.check_output_with_option(use_gpu, flatten=True)
            self.assertTrue(
                PassVersionChecker.IsCompatible('tensorrt_subgraph_pass'))


if __name__ == "__main__":
    unittest.main()

[((9538, 9553), 'unittest.main', 'unittest.main', ([], {}), '()\n', (9551, 9553), False, 'import unittest\n'), ((1836, 1864), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (1862, 1864), True, 'import paddle.fluid.core as core\n'), ((3028, 3056), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (3054, 3056), True, 'import paddle.fluid.core as core\n'), ((4208, 4236), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (4234, 4236), True, 'import paddle.fluid.core as core\n'), ((5358, 5386), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (5384, 5386), True, 'import paddle.fluid.core as core\n'), ((6361, 6389), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (6387, 6389), True, 'import paddle.fluid.core as core\n'), ((7340, 7368), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (7366, 7368), True, 'import paddle.fluid.core as core\n'), ((8310, 8338), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (8336, 8338), True, 'import paddle.fluid.core as core\n'), ((9279, 9307), 'paddle.fluid.core.is_compiled_with_cuda', 'core.is_compiled_with_cuda', ([], {}), '()\n', (9305, 9307), True, 'import paddle.fluid.core as core\n'), ((971, 1031), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (990, 1031), True, 'import paddle.fluid as fluid\n'), ((1052, 1115), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[-1, 3, -1, -1]', 'dtype': '"""float32"""'}), "(name='data', shape=[-1, 3, -1, -1], dtype='float32')\n", (1062, 1115), True, 'import paddle.fluid as fluid\n'), ((1159, 1217), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'dim': '[2, -1]', 'keep_dim': '(True)'}), '(data, dim=[2, -1], keep_dim=True)\n', (1183, 1217), True, 'import paddle.fluid as fluid\n'), ((1253, 1303), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (1276, 1303), True, 'import paddle.fluid as fluid\n'), ((2152, 2212), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (2171, 2212), True, 'import paddle.fluid as fluid\n'), ((2233, 2296), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[-1, 3, -1, -1]', 'dtype': '"""float32"""'}), "(name='data', shape=[-1, 3, -1, -1], dtype='float32')\n", (2243, 2296), True, 'import paddle.fluid as fluid\n'), ((2340, 2385), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'keep_dim': '(True)'}), '(data, keep_dim=True)\n', (2364, 2385), True, 'import paddle.fluid as fluid\n'), ((2404, 2454), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (2427, 2454), True, 'import paddle.fluid as fluid\n'), ((3338, 3398), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (3357, 3398), True, 'import paddle.fluid as fluid\n'), ((3419, 3482), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[-1, 3, -1, -1]', 'dtype': '"""float32"""'}), "(name='data', shape=[-1, 3, -1, -1], dtype='float32')\n", (3429, 3482), True, 'import paddle.fluid as fluid\n'), ((3526, 3584), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'dim': '[2, -1]', 'keep_dim': '(True)'}), '(data, dim=[2, -1], keep_dim=True)\n', (3550, 3584), True, 'import paddle.fluid as fluid\n'), ((3620, 3670), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (3643, 3670), True, 'import paddle.fluid as fluid\n'), ((4517, 4577), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (4536, 4577), True, 'import paddle.fluid as fluid\n'), ((4598, 4661), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[-1, 3, 56, 56]', 'dtype': '"""float32"""'}), "(name='data', shape=[-1, 3, 56, 56], dtype='float32')\n", (4608, 4661), True, 'import paddle.fluid as fluid\n'), ((4705, 4750), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'keep_dim': '(True)'}), '(data, keep_dim=True)\n', (4729, 4750), True, 'import paddle.fluid as fluid\n'), ((4769, 4819), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (4792, 4819), True, 'import paddle.fluid as fluid\n'), ((5670, 5730), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (5689, 5730), True, 'import paddle.fluid as fluid\n'), ((5751, 5813), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[3, 3, 56, 56]', 'dtype': '"""float32"""'}), "(name='data', shape=[3, 3, 56, 56], dtype='float32')\n", (5761, 5813), True, 'import paddle.fluid as fluid\n'), ((5857, 5915), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'dim': '[2, -1]', 'keep_dim': '(True)'}), '(data, dim=[2, -1], keep_dim=True)\n', (5881, 5915), True, 'import paddle.fluid as fluid\n'), ((5951, 6001), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (5974, 6001), True, 'import paddle.fluid as fluid\n'), ((6676, 6736), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (6695, 6736), True, 'import paddle.fluid as fluid\n'), ((6757, 6819), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[4, 3, 56, 56]', 'dtype': '"""float32"""'}), "(name='data', shape=[4, 3, 56, 56], dtype='float32')\n", (6767, 6819), True, 'import paddle.fluid as fluid\n'), ((6863, 6908), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'keep_dim': '(True)'}), '(data, keep_dim=True)\n', (6887, 6908), True, 'import paddle.fluid as fluid\n'), ((6927, 6977), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (6950, 6977), True, 'import paddle.fluid as fluid\n'), ((7652, 7712), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (7671, 7712), True, 'import paddle.fluid as fluid\n'), ((7733, 7795), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[4, 3, 56, 56]', 'dtype': '"""float32"""'}), "(name='data', shape=[4, 3, 56, 56], dtype='float32')\n", (7743, 7795), True, 'import paddle.fluid as fluid\n'), ((7839, 7884), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'keep_dim': '(True)'}), '(data, keep_dim=True)\n', (7863, 7884), True, 'import paddle.fluid as fluid\n'), ((7903, 7953), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (7926, 7953), True, 'import paddle.fluid as fluid\n'), ((8622, 8682), 'paddle.fluid.program_guard', 'fluid.program_guard', (['self.main_program', 'self.startup_program'], {}), '(self.main_program, self.startup_program)\n', (8641, 8682), True, 'import paddle.fluid as fluid\n'), ((8703, 8765), 'paddle.fluid.data', 'fluid.data', ([], {'name': '"""data"""', 'shape': '[4, 3, 56, 56]', 'dtype': '"""float32"""'}), "(name='data', shape=[4, 3, 56, 56], dtype='float32')\n", (8713, 8765), True, 'import paddle.fluid as fluid\n'), ((8809, 8854), 'paddle.fluid.layers.reduce_mean', 'fluid.layers.reduce_mean', (['data'], {'keep_dim': '(True)'}), '(data, keep_dim=True)\n', (8833, 8854), True, 'import paddle.fluid as fluid\n'), ((8873, 8923), 'paddle.fluid.layers.batch_norm', 'fluid.layers.batch_norm', (['reduce_mean'], {'is_test': '(True)'}), '(reduce_mean, is_test=True)\n', (8896, 8923), True, 'import paddle.fluid as fluid\n'), ((2003, 2060), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (2034, 2060), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((3195, 3252), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (3226, 3252), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((4375, 4432), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (4406, 4432), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((5525, 5582), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (5556, 5582), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((6528, 6585), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (6559, 6585), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((7507, 7564), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (7538, 7564), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((8477, 8534), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (8508, 8534), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((9446, 9503), 'paddle.fluid.core.PassVersionChecker.IsCompatible', 'PassVersionChecker.IsCompatible', (['"""tensorrt_subgraph_pass"""'], {}), "('tensorrt_subgraph_pass')\n", (9477, 9503), False, 'from paddle.fluid.core import PassVersionChecker\n'), ((1348, 1380), 'numpy.random.random', 'np.random.random', (['[3, 3, 56, 56]'], {}), '([3, 3, 56, 56])\n', (1364, 1380), True, 'import numpy as np\n'), ((2499, 2531), 'numpy.random.random', 'np.random.random', (['[3, 3, 56, 56]'], {}), '([3, 3, 56, 56])\n', (2515, 2531), True, 'import numpy as np\n'), ((3715, 3747), 'numpy.random.random', 'np.random.random', (['[3, 3, 56, 56]'], {}), '([3, 3, 56, 56])\n', (3731, 3747), True, 'import numpy as np\n'), ((4864, 4896), 'numpy.random.random', 'np.random.random', (['[3, 3, 56, 56]'], {}), '([3, 3, 56, 56])\n', (4880, 4896), True, 'import numpy as np\n'), ((6046, 6078), 'numpy.random.random', 'np.random.random', (['[3, 3, 56, 56]'], {}), '([3, 3, 56, 56])\n', (6062, 6078), True, 'import numpy as np\n'), ((7022, 7054), 'numpy.random.random', 'np.random.random', (['[4, 3, 56, 56]'], {}), '([4, 3, 56, 56])\n', (7038, 7054), True, 'import numpy as np\n'), ((7998, 8030), 'numpy.random.random', 'np.random.random', (['[4, 3, 56, 56]'], {}), '([4, 3, 56, 56])\n', (8014, 8030), True, 'import numpy as np\n'), ((8968, 9000), 'numpy.random.random', 'np.random.random', (['[4, 3, 56, 56]'], {}), '([4, 3, 56, 56])\n', (8984, 9000), True, 'import numpy as np\n')]