Files
paddlepaddle--paddle/test/ir/inference/test_save_optimized_model_pass.py
2026-07-13 12:40:42 +08:00

238 lines
7.8 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 tempfile
import unittest
import numpy as np
import paddle
from paddle.inference import Config, PrecisionType, create_predictor
from paddle.jit import to_static
from paddle.static import InputSpec
from paddle.vision.models import alexnet
class TestSaveOptimizedModelPass:
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
net = alexnet(True)
model = to_static(
net,
input_spec=[InputSpec(shape=[None, 3, 224, 224], name='x')],
full_graph=True,
)
paddle.jit.save(
model, os.path.join(self.temp_dir.name, 'alexnet/inference')
)
def tearDown(self):
self.temp_dir.cleanup()
def get_baseline(self):
predictor = self.init_predictor(save_optimized_model=True)
inputs = [
paddle.to_tensor(0.1 * np.ones([1, 3, 224, 224]).astype(np.float32))
]
outputs = predictor.run(inputs)
return outputs[0]
def get_test_output(self):
predictor = self.init_predictor(save_optimized_model=False)
inputs = [
paddle.to_tensor(0.1 * np.ones([1, 3, 224, 224]).astype(np.float32))
]
outputs = predictor.run(inputs)
return outputs[0]
def test_output(self):
if paddle.is_compiled_with_cuda():
baseline = self.get_baseline()
test_output = self.get_test_output()
np.testing.assert_allclose(
baseline.numpy().flatten(),
test_output.numpy().flatten(),
)
class TestSaveOptimizedModelPassWithGPU(
TestSaveOptimizedModelPass, unittest.TestCase
):
def init_predictor(self, save_optimized_model: bool):
if save_optimized_model is True:
config = Config(
os.path.join(self.temp_dir.name, 'alexnet/inference.json'),
os.path.join(self.temp_dir.name, 'alexnet/inference.pdiparams'),
)
config.enable_use_gpu(256, 0, PrecisionType.Half)
config.enable_memory_optim()
config.switch_ir_optim(True)
config.set_optim_cache_dir(
os.path.join(self.temp_dir.name, 'alexnet')
)
config.enable_save_optim_model(True)
else:
config = Config(
os.path.join(self.temp_dir.name, 'alexnet/_optimized.json'),
os.path.join(
self.temp_dir.name, 'alexnet/_optimized.pdiparams'
),
)
config.enable_use_gpu(256, 0, PrecisionType.Half)
config.enable_memory_optim()
config.switch_ir_optim(False)
predictor = create_predictor(config)
return predictor
class TestSavePirOptimizedModelPassWithGPU(unittest.TestCase):
def setUp(self):
self.temp_dir = "./"
net = alexnet(True)
with paddle.pir_utils.DygraphPirGuard():
model = to_static(
net,
input_spec=[InputSpec(shape=[None, 3, 224, 224], name='x')],
full_graph=True,
)
paddle.jit.save(
model, os.path.join(self.temp_dir, 'alexnet/inference')
)
def tearDown(self):
# self.temp_dir.cleanup()
pass
def get_baseline(self):
predictor = self.init_predictor(save_optimized_model=True)
inputs = [
paddle.to_tensor(0.1 * np.ones([1, 3, 224, 224]).astype(np.float32))
]
outputs = predictor.run(inputs)
return outputs[0]
def get_test_output(self):
predictor = self.init_predictor(save_optimized_model=False)
inputs = [
paddle.to_tensor(0.1 * np.ones([1, 3, 224, 224]).astype(np.float32))
]
outputs = predictor.run(inputs)
return outputs[0]
def test_output(self):
if paddle.is_compiled_with_cuda():
baseline = self.get_baseline()
test_output = self.get_test_output()
np.testing.assert_allclose(
baseline.numpy().flatten(),
test_output.numpy().flatten(),
)
def init_predictor(self, save_optimized_model: bool):
if save_optimized_model is True:
config = Config(
os.path.join(self.temp_dir, 'alexnet/inference.json'),
os.path.join(self.temp_dir, 'alexnet/inference.pdiparams'),
)
config.enable_use_gpu(256, 0, PrecisionType.Half)
config.enable_memory_optim()
config.switch_ir_optim(True)
config.set_optim_cache_dir(os.path.join(self.temp_dir, 'alexnet'))
config.enable_save_optim_model(True)
else:
config = Config(
os.path.join(self.temp_dir, 'alexnet/_optimized.json'),
os.path.join(self.temp_dir, 'alexnet/_optimized.pdiparams'),
)
config.enable_use_gpu(256, 0, PrecisionType.Half)
config.enable_memory_optim()
config.switch_ir_optim(False)
config.enable_new_executor()
config.enable_new_ir()
predictor = create_predictor(config)
return predictor
class TestSaveOptimizedModelPassWithTRT(
TestSaveOptimizedModelPass, unittest.TestCase
):
def init_predictor(self, save_optimized_model: bool):
if save_optimized_model is True:
config = Config(
os.path.join(self.temp_dir.name, 'alexnet/inference.json'),
os.path.join(self.temp_dir.name, 'alexnet/inference.pdiparams'),
)
config.enable_use_gpu(256, 0)
config.enable_tensorrt_engine(
workspace_size=1 << 30,
max_batch_size=1,
min_subgraph_size=3,
precision_mode=PrecisionType.Half,
use_static=True,
use_calib_mode=False,
)
config.set_trt_dynamic_shape_info(
{"x": [1, 3, 224, 224], "flatten_0.tmp_0": [1, 9216]},
{"x": [1, 3, 224, 224], "flatten_0.tmp_0": [1, 9216]},
{"x": [1, 3, 224, 224], "flatten_0.tmp_0": [1, 9216]},
)
config.exp_disable_tensorrt_ops(["flatten_contiguous_range"])
config.enable_memory_optim()
config.switch_ir_optim(True)
config.set_optim_cache_dir(
os.path.join(self.temp_dir.name, 'alexnet')
)
config.enable_save_optim_model(True)
else:
config = Config(
os.path.join(self.temp_dir.name, 'alexnet/_optimized.json'),
os.path.join(
self.temp_dir.name, 'alexnet/_optimized.pdiparams'
),
)
config.enable_use_gpu(256, 0)
config.enable_tensorrt_engine(
workspace_size=1 << 30,
max_batch_size=1,
min_subgraph_size=3,
precision_mode=PrecisionType.Half,
use_static=True,
use_calib_mode=False,
)
config.enable_memory_optim()
config.switch_ir_optim(False)
predictor = create_predictor(config)
return predictor
if __name__ == '__main__':
unittest.main()