Files
paddlepaddle--paddle/test/tensorrt/test_export.py
T
2026-07-13 12:40:42 +08:00

365 lines
12 KiB
Python

# Copyright (c) 2024 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 tempfile
import unittest
import numpy as np
import paddle
import paddle.inference as paddle_infer
import paddle.nn.functional as F
from paddle import Tensor, nn
from paddle.static import InputSpec
from paddle.tensorrt.export import (
Input,
TensorRTConfig,
_convert_,
)
from paddle.tensorrt.util import (
predict_program,
)
class LeNetMultiInput(nn.Layer):
"""LeNet model modified to accept two inputs."""
def __init__(self, num_classes: int = 10) -> None:
super().__init__()
self.num_classes = num_classes
# Convolution layers for the first input
self.features1 = nn.Sequential(
nn.Conv2D(1, 6, 3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2D(2, 2),
nn.Conv2D(6, 16, 5, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2D(2, 2),
)
# Convolution layers for the second input
self.features2 = nn.Sequential(
nn.Conv2D(1, 6, 3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2D(2, 2),
nn.Conv2D(6, 16, 5, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2D(2, 2),
)
# Fully connected layers
if num_classes > 0:
self.fc = nn.Sequential(
nn.Linear(400 * 2, 120), # Adjusted for two inputs
nn.Linear(120, 84),
nn.Linear(84, num_classes),
)
def forward(self, input1: Tensor, input2: Tensor) -> Tensor:
# Apply feature extraction on both inputs
x1 = self.features1(input1)
x2 = self.features2(input2)
# Flatten both feature maps
x1 = paddle.flatten(x1, 1)
x2 = paddle.flatten(x2, 1)
# Concatenate the features from both inputs
x = paddle.concat([x1, x2], axis=1)
if self.num_classes > 0:
x = self.fc(x)
return x
class CumsumModel(nn.Layer):
def __init__(self, input_dim):
super().__init__()
self.linear = nn.Linear(input_dim, input_dim)
def forward(self, x):
linear_out = self.linear(x)
relu_out = F.relu(linear_out)
axis = paddle.full([1], 2, dtype='int64')
out = paddle.cumsum(relu_out, axis=axis)
return out
class TestConvert(unittest.TestCase):
def setUp(self):
paddle.seed(2024)
self.temp_dir = tempfile.TemporaryDirectory()
self.save_path = os.path.join(self.temp_dir.name, 'tensor_axis_cumsum')
self.place = (
paddle.CUDAPlace(0)
if paddle.is_compiled_with_cuda()
else paddle.CPUPlace()
)
def test_paddle_to_tensorrt_conversion_cumsum(self):
paddle.enable_static()
np_x = np.random.randn(9, 10, 11).astype('float32')
with paddle.pir_utils.IrGuard():
main_prog = paddle.static.Program()
startup_prog = paddle.static.Program()
with paddle.static.program_guard(main_prog, startup_prog):
x = paddle.static.data(
shape=np_x.shape, name='x', dtype=np_x.dtype
)
model = CumsumModel(input_dim=np_x.shape[-1])
out = model(x)
loss = paddle.mean(out)
sgd = paddle.optimizer.SGD(learning_rate=0.0)
sgd.minimize(paddle.mean(out))
exe = paddle.static.Executor(self.place)
exe.run(startup_prog)
static_out = exe.run(feed={'x': np_x}, fetch_list=[out])
# run infer
paddle.static.save_inference_model(
self.save_path, [x], [out], exe
)
config = paddle_infer.Config(
self.save_path + '.json', self.save_path + '.pdiparams'
)
config.enable_new_ir()
config.enable_new_executor()
config.use_optimized_model(True)
# Set input
input_config = Input(
min_input_shape=(9, 10, 11),
optim_input_shape=(9, 10, 11),
max_input_shape=(9, 10, 11),
)
# Create a TensorRTConfig with inputs as a required field.
trt_config = TensorRTConfig(inputs=[input_config])
trt_save_path = os.path.join(self.temp_dir.name, 'trt')
trt_config.save_model_dir = trt_save_path
trt_config.refit_params_path = self.save_path + '.pdiparams'
model_dir = self.save_path
# Obtain tensorrt_engine_op by passing the model path and trt_config.(converted_program)
program_with_trt = paddle.tensorrt.convert(model_dir, trt_config)
# Create a config for inference.
config = paddle_infer.Config(
trt_config.save_model_dir + '.json',
trt_config.save_model_dir + '.pdiparams',
)
if paddle.is_compiled_with_cuda():
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
predictor = paddle_infer.create_predictor(config)
paddle.disable_static()
for i, input_instance in enumerate(trt_config.inputs):
min_data, _, max_data = input_instance.generate_input_data()
model_inputs = paddle.to_tensor(min_data)
output_converted = predictor.run([model_inputs])
class TestConvert_(unittest.TestCase):
def test_run(self):
with paddle.pir_utils.IrGuard():
input_config = Input(
min_input_shape=(9, 10, 11),
optim_input_shape=(9, 10, 11),
max_input_shape=(10, 10, 11),
)
trt_config = TensorRTConfig(inputs=[input_config])
for i, input_instance in enumerate(trt_config.inputs):
min_data, _, max_data = input_instance.generate_input_data()
paddle.disable_static()
x = paddle.to_tensor(min_data)
net = CumsumModel(input_dim=min_data.shape[-1])
out = net(x)
input_spec = [
InputSpec(shape=[None, 10, 11], dtype='float32', name='x')
]
program_with_trt, scope = _convert_(
net,
input_spec=input_spec,
config=trt_config,
)
output_var = program_with_trt.list_vars()[-1]
output_converted = predict_program(
program_with_trt,
{"x": min_data},
[output_var],
scope=scope,
)
output_expected = out.numpy()
output_converted_np = output_converted[0]
# Check that the results are close to each other within a tolerance of 1e-2
np.testing.assert_allclose(
output_expected,
output_converted_np,
rtol=1e-2,
atol=1e-2,
err_msg="Outputs are not within the 1e-2 tolerance",
)
class TestConvertMultipleInputs(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
self.save_path = os.path.join(
self.temp_dir.name, 'tensor_axis_cumsum_multiple'
)
self.place = (
paddle.CUDAPlace(0)
if paddle.is_compiled_with_cuda()
else paddle.CPUPlace()
)
def test_run(self):
with paddle.pir_utils.IrGuard():
input_config = Input(
min_input_shape=(1, 1, 28, 28),
optim_input_shape=(1, 1, 28, 28),
max_input_shape=(1, 1, 28, 28),
)
input_config2 = Input(
min_input_shape=(1, 1, 28, 28),
optim_input_shape=(1, 1, 28, 28),
max_input_shape=(1, 1, 28, 28),
)
trt_config = TensorRTConfig(inputs=[input_config, input_config2])
trt_config.save_model_dir = os.path.join(self.temp_dir.name, 'trt')
min_data_list = []
max_data_list = []
for i, input_instance in enumerate(trt_config.inputs):
min_data, _, max_data = input_instance.generate_input_data()
min_data_list.append(min_data)
max_data_list.append(max_data)
paddle.disable_static()
x = [paddle.to_tensor(md) for md in min_data_list]
net = LeNetMultiInput()
out = net(*x)
input_spec = [
InputSpec(
shape=min_data_list[0].shape, dtype='float32', name='input1'
),
InputSpec(
shape=min_data_list[1].shape, dtype='float32', name='input2'
),
]
program_with_trt, scope = _convert_(
net,
input_spec=input_spec,
config=trt_config,
full_graph=True,
)
config = paddle_infer.Config(
trt_config.save_model_dir + '.json',
trt_config.save_model_dir + '.pdiparams',
)
if paddle.is_compiled_with_cuda():
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
predictor = paddle_infer.create_predictor(config)
output_converted = predictor.run(x)
output_converted_np = output_converted[0]
output_expected = out.numpy()
np.testing.assert_allclose(
output_expected,
output_converted_np,
rtol=1e-2,
atol=1e-2,
err_msg="Outputs are not within the 1e-2 tolerance",
)
class TestConvertPredictor(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
self.save_path = os.path.join(self.temp_dir.name, 'tensor_axis_cumsum')
self.place = (
paddle.CUDAPlace(0)
if paddle.is_compiled_with_cuda()
else paddle.CPUPlace()
)
def test_run(self):
input_config = Input(
min_input_shape=(9, 10, 11),
optim_input_shape=(9, 10, 11),
max_input_shape=(10, 10, 11),
)
trt_config = TensorRTConfig(inputs=[input_config])
trt_config.save_model_dir = os.path.join(self.temp_dir.name, 'trt')
min_data, _, max_data = input_config.generate_input_data()
net = CumsumModel(input_dim=min_data.shape[-1])
x = paddle.to_tensor(min_data)
out = net(x).numpy()
input_spec = [
InputSpec(shape=[None, 10, 11], dtype='float32', name='x')
]
program_with_trt, scope = _convert_(
net,
input_spec=input_spec,
config=trt_config,
)
config = paddle_infer.Config(
trt_config.save_model_dir + '.json',
trt_config.save_model_dir + '.pdiparams',
)
if paddle.is_compiled_with_cuda():
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
predictor = paddle_infer.create_predictor(config)
output_converted = predictor.run([x])
output_converted_np = output_converted[0]
np.testing.assert_allclose(
out,
output_converted_np,
rtol=1e-2,
atol=1e-2,
err_msg="Outputs are not within the 1e-2 tolerance",
)
if __name__ == "__main__":
unittest.main()