paddlenlp

1
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import os
16

17
import paddle
18

19
from .log import logger
20

21
__all__ = ["export_inference_model"]
22

23

24
def _prune_input_spec(input_spec, program, targets):
25
    # try to prune static program to figure out pruned input spec
26
    # so we perform following operations in static mode
27
    device = paddle.get_device()
28
    paddle.enable_static()
29
    paddle.set_device(device)
30
    pruned_input_spec = []
31
    program = program.clone()
32
    program = program._prune(targets=targets)
33
    global_block = program.global_block()
34
    for spec in input_spec:
35
        try:
36
            global_block.var(spec.name)
37
            pruned_input_spec.append(spec)
38
        except Exception:
39
            pass
40
    paddle.disable_static(place=device)
41
    return pruned_input_spec
42

43

44
def export_inference_model(
45
    model,
46
    input_spec,
47
    save_dir="./output",
48
    save_name="model",
49
    export_quant_model=False,
50
    quanter=None,
51
):
52
    if not os.path.exists(save_dir):
53
        os.makedirs(save_dir)
54

55
    static_model = paddle.jit.to_static(model, input_spec)
56
    pruned_input_spec = _prune_input_spec(input_spec, static_model.forward.main_program, static_model.forward.outputs)
57

58
    if export_quant_model:
59
        quanter.save_quantized_model(model, os.path.join(save_dir, save_name), input_spec=pruned_input_spec)
60
        logger.info("export quantized inference model saved in {}".format(save_dir))
61
        return
62

63
    paddle.jit.save(static_model, os.path.join(save_dir, save_name), input_spec=pruned_input_spec)
64
    logger.info("export inference model saved in {}".format(save_dir))
65