You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
PaddleSpeech/paddlespeech/server/utils/paddle_predictor.py

99 lines
3.2 KiB

# 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.
import os
from typing import List
from typing import Optional
import paddle
from paddle.inference import Config
from paddle.inference import create_predictor
def init_predictor(model_dir: Optional[os.PathLike]=None,
model_file: Optional[os.PathLike]=None,
params_file: Optional[os.PathLike]=None,
predictor_conf: dict=None):
"""Create predictor with Paddle inference
Args:
model_dir (Optional[os.PathLike], optional): The path of the static model saved in the model layer. Defaults to None.
model_file (Optional[os.PathLike], optional): *.pdmodel file path. Defaults to None.
params_file (Optional[os.PathLike], optional): *.pdiparams file path.. Defaults to None.
predictor_conf (dict, optional): The configuration parameters of predictor. Defaults to None.
Returns:
predictor (PaddleInferPredictor): created predictor
"""
if model_dir is not None:
assert os.path.isdir(model_dir), 'Please check model dir.'
config = Config(args.model_dir)
else:
assert os.path.isfile(model_file) and os.path.isfile(
params_file), 'Please check model and parameter files.'
config = Config(model_file, params_file)
# set device
if predictor_conf["device"]:
device = predictor_conf["device"]
else:
device = paddle.get_device()
if "gpu" in device:
gpu_id = device.split(":")[-1]
config.enable_use_gpu(1000, int(gpu_id))
# IR optim
if predictor_conf["switch_ir_optim"]:
config.switch_ir_optim()
# glog
if not predictor_conf["glog_info"]:
config.disable_glog_info()
# config summary
if predictor_conf["summary"]:
print(config.summary())
# memory optim
config.enable_memory_optim()
predictor = create_predictor(config)
return predictor
def run_model(predictor, input: List) -> List:
""" run predictor
Args:
predictor: paddle inference predictor
input (list): The input of predictor
Returns:
list: result list
"""
input_names = predictor.get_input_names()
for i, name in enumerate(input_names):
input_handle = predictor.get_input_handle(name)
input_handle.copy_from_cpu(input[i])
# do the inference
predictor.run()
results = []
# get out data from output tensor
output_names = predictor.get_output_names()
for i, name in enumerate(output_names):
output_handle = predictor.get_output_handle(name)
output_data = output_handle.copy_to_cpu()
results.append(output_data)
return results