chore: import upstream snapshot with attribution

test/ir/inference/test_trt_explicit_quantization_model.py

@@ -0,0 +1,175 @@
+# Copyright (c) 2023 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.
+
+import os
+import shutil
+import tempfile
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+from paddle.base.executor import global_scope
+from paddle.base.framework import IrGraph
+from paddle.inference import Config, PrecisionType, create_predictor
+from paddle.static.quantization import QuantizationTransformPassV2
+
+
+class TestExplicitQuantizationModel:
+    def setUp(self):
+        paddle.enable_static()
+        np.random.seed(1024)
+        paddle.seed(1024)
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.path = os.path.join(
+            self.temp_dir.name, 'trt_explicit', self.__class__.__name__
+        )
+
+    def tearDown(self):
+        shutil.rmtree(self.path)
+
+    def build_program(self):
+        train_prog = paddle.static.Program()
+        with paddle.static.program_guard(train_prog):
+            image = paddle.static.data(
+                name='image', shape=[None, 1, 28, 28], dtype='float32'
+            )
+            label = paddle.static.data(
+                name='label', shape=[None, 1], dtype='int64'
+            )
+            model = self.build_model()
+            out = model.net(input=image, class_dim=10)
+            cost = paddle.nn.functional.loss.cross_entropy(
+                input=out, label=label
+            )
+            avg_cost = paddle.mean(x=cost)
+            acc_top1 = paddle.metric.accuracy(input=out, label=label, k=1)
+            optimizer = paddle.optimizer.Momentum(
+                momentum=0.9,
+                learning_rate=0.01,
+                weight_decay=paddle.regularizer.L2Decay(4e-5),
+            )
+            optimizer.minimize(avg_cost)
+        val_prog = train_prog.clone(for_test=True)
+        place = paddle.CUDAPlace(0)
+        exe = paddle.static.Executor(place)
+        exe.run(paddle.static.default_startup_program())
+
+        def transform(x):
+            return np.reshape(x, [1, 28, 28]) - 127.5 / 127.5
+
+        train_dataset = paddle.vision.datasets.MNIST(
+            mode='train', backend='cv2', transform=transform
+        )
+        train_loader = paddle.io.DataLoader(
+            train_dataset,
+            places=place,
+            feed_list=[image, label],
+            drop_last=True,
+            return_list=False,
+            batch_size=64,
+        )
+
+        def train(program, stop_iter=128):
+            for it, data in enumerate(train_loader):
+                if it == 0:
+                    self.input_data = data[0]['image']
+                loss, top1 = exe.run(
+                    program, feed=data, fetch_list=[avg_cost, acc_top1]
+                )
+                scope = global_scope()
+                if it == stop_iter:
+                    break
+
+        train(train_prog)
+
+        scope = global_scope()
+
+        def insert_qdq(program, scope, place, for_test=False):
+            graph = IrGraph(core.Graph(program.desc), for_test=for_test)
+            transform_pass = QuantizationTransformPassV2(
+                scope=scope,
+                place=place,
+                activation_quantize_type='moving_average_abs_max',
+                weight_quantize_type='channel_wise_abs_max',
+            )
+            transform_pass.apply(graph)
+            quant_program = graph.to_program()
+            return quant_program
+
+        quant_train_prog = insert_qdq(train_prog, scope, place, for_test=False)
+        quant_val_prog = insert_qdq(val_prog, scope, place, for_test=True)
+
+        train(quant_train_prog)
+
+        path_prefix = os.path.join(self.path, 'inference')
+        paddle.static.save_inference_model(
+            path_prefix, [image], [out], exe, program=quant_val_prog
+        )
+
+    def infer_program(self, trt_int8=False, collect_shape=False):
+        config = Config(
+            os.path.join(self.path, 'inference.pdmodel'),
+            os.path.join(self.path, 'inference.pdiparams'),
+        )
+        config.enable_use_gpu(256, 0, PrecisionType.Float32)
+        config.enable_memory_optim()
+        if trt_int8:
+            precision_mode = PrecisionType.Int8
+        else:
+            precision_mode = PrecisionType.Float32
+        config.enable_tensorrt_engine(
+            workspace_size=1 << 30,
+            max_batch_size=1,
+            min_subgraph_size=3,
+            precision_mode=precision_mode,
+            use_static=False,
+            use_calib_mode=False,
+        )
+        if trt_int8:
+            config.enable_tensorrt_explicit_quantization()
+        shape_path = self.path + ".shape.txt"
+        if collect_shape:
+            config.collect_shape_range_info(shape_path)
+        else:
+            config.enable_tuned_tensorrt_dynamic_shape(shape_path)
+        config.disable_glog_info()
+        predictor = create_predictor(config)
+        input_names = predictor.get_input_names()
+        input_tensor = predictor.get_input_handle(input_names[0])
+        input_tensor.reshape(self.input_data.shape())
+        input_tensor.share_external_data(self.input_data)
+        predictor.run()
+        output_names = predictor.get_output_names()
+        output_tensor = predictor.get_output_handle(output_names[0])
+        output_data = output_tensor.copy_to_cpu()
+        return output_data
+
+    def test_model(self):
+        with paddle.pir_utils.OldIrGuard():
+            self.build_program()
+            self.infer_program(trt_int8=False, collect_shape=True)
+            baseline_output = self.infer_program(
+                trt_int8=False, collect_shape=False
+            )
+            self.infer_program(trt_int8=True, collect_shape=True)
+            trt_output = self.infer_program(trt_int8=True, collect_shape=False)
+            trt_predict = np.argmax(trt_output, axis=1)
+            baseline_predict = np.argmax(baseline_output, axis=1)
+            same = (trt_predict == baseline_predict).sum() / len(trt_predict)
+            self.assertGreaterEqual(
+                same,
+                0.9,
+                "There are more then 10% output difference between int8 and float32 inference.",
+            )