wehub-resource-sync
133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
# 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 site
|
|
import sys
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from utils import check_output_allclose
|
|
|
|
import paddle
|
|
from paddle.utils.cpp_extension.extension_utils import run_cmd
|
|
|
|
|
|
class GapTestNet(paddle.nn.Layer):
|
|
def __init__(self, gap_op):
|
|
super().__init__()
|
|
self.test_attr1 = [1, 2, 3]
|
|
self.test_attr2 = 1
|
|
self.linear = paddle.nn.Linear(96, 1)
|
|
self.conv1 = paddle.nn.Conv2D(3, 6, kernel_size=3)
|
|
self.conv2 = paddle.nn.Conv2D(6, 3, kernel_size=3)
|
|
self.gap = gap_op
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.conv2(x)
|
|
x = self.gap(x, self.test_attr1, self.test_attr2)
|
|
x = paddle.flatten(x)
|
|
x = self.linear(x)
|
|
return x
|
|
|
|
|
|
class TestNewCustomOpSetUpInstall(unittest.TestCase):
|
|
def setUp(self):
|
|
# TODO(ming1753): skip window CI because run_cmd(cmd) filed
|
|
if os.name != 'nt':
|
|
cur_dir = os.path.dirname(os.path.abspath(__file__))
|
|
# compile, install the custom op egg into site-packages under background
|
|
cmd = f'cd {cur_dir} && {sys.executable} inference_gap_setup.py install'
|
|
run_cmd(cmd)
|
|
|
|
site_dir = site.getsitepackages()[0]
|
|
|
|
custom_install_path = [
|
|
x for x in os.listdir(site_dir) if 'gap_op_setup' in x
|
|
]
|
|
|
|
assert len(custom_install_path) == 2, (
|
|
f"Matched egg number is {len(custom_install_path)}."
|
|
)
|
|
|
|
sys.path.append(os.path.join(site_dir, custom_install_path[0]))
|
|
|
|
# usage: import the package directly
|
|
import gap_op_setup
|
|
|
|
# `custom_relu_dup` is same as `custom_relu_dup`
|
|
self.custom_op = gap_op_setup.gap
|
|
|
|
# config seed
|
|
SEED = 2021
|
|
paddle.seed(SEED)
|
|
paddle.framework.random._manual_program_seed(SEED)
|
|
|
|
def test_all(self):
|
|
if paddle.is_compiled_with_cuda() and os.name != 'nt':
|
|
self._test_static_save_and_run_inference_predictor()
|
|
|
|
def _test_static_save_and_run_inference_predictor(self):
|
|
np_data = np.ones((32, 3, 7, 7)).astype("float32")
|
|
path_prefix = "custom_op_inference/inference_gap_op"
|
|
model = GapTestNet(self.custom_op)
|
|
x = paddle.to_tensor(np_data)
|
|
y = model(x)
|
|
paddle.jit.save(
|
|
model,
|
|
path_prefix,
|
|
input_spec=[
|
|
paddle.static.InputSpec(shape=[32, 3, 7, 7], dtype='float32')
|
|
],
|
|
)
|
|
|
|
from paddle.inference import Config, create_predictor
|
|
|
|
# load inference model
|
|
config = Config(path_prefix + ".pdmodel", path_prefix + ".pdiparams")
|
|
config.enable_use_gpu(500, 0)
|
|
config.enable_tensorrt_engine(
|
|
workspace_size=1 << 30,
|
|
max_batch_size=1,
|
|
min_subgraph_size=0,
|
|
precision_mode=paddle.inference.PrecisionType.Float32,
|
|
use_static=True,
|
|
use_calib_mode=False,
|
|
)
|
|
config.set_trt_dynamic_shape_info(
|
|
{"x": [32, 3, 7, 7]},
|
|
{"x": [32, 3, 7, 7]},
|
|
{"x": [32, 3, 7, 7]},
|
|
)
|
|
predictor = create_predictor(config)
|
|
input_tensor = predictor.get_input_handle(
|
|
predictor.get_input_names()[0]
|
|
)
|
|
input_tensor.reshape(np_data.shape)
|
|
input_tensor.copy_from_cpu(np_data.copy())
|
|
predictor.run()
|
|
output_tensor = predictor.get_output_handle(
|
|
predictor.get_output_names()[0]
|
|
)
|
|
predict_infer = output_tensor.copy_to_cpu()
|
|
predict = y.numpy().flatten()
|
|
|
|
predict_infer = np.array(predict_infer).flatten()
|
|
check_output_allclose(predict, predict_infer, "predict")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|