[ConvBert] support static export

[JunnYu]

PR types

Bug fixes

PR changes

Models

Description

支持convbert静态图导出。

• 使用paddle.shape而不是tensor.shape获取tensor的形状。
• F.unfold无法在变长输入情况下使用，因此采用for循环加slice方式处理。

导出代码

from paddlenlp.transformers import ConvBertModel, ConvBertTokenizer
import paddle
import os
import argparse


def get_args(add_help=True):
    """get_args
    Parse all args using argparse lib
    Args:
        add_help: Whether to add -h option on args
    Returns:
        An object which contains many parameters used for inference.
    """
    parser = argparse.ArgumentParser(
        description='Paddlenlp Classification Training', add_help=add_help)
    parser.add_argument(
        "--model_path",
        default="convbert-base",
        type=str,
        help="Path of the trained model to be exported.", )
    parser.add_argument(
        '--save_inference_dir',
        default='./convbert_infer',
        help='path where to save')

    args = parser.parse_args()
    return args


def export(args):
    # build model
    model = ConvBertModel.from_pretrained(args.model_path)
    tokenizer = ConvBertTokenizer.from_pretrained(args.model_path)
    model.eval()

    # decorate model with jit.save
    model = paddle.jit.to_static(
        model,
        input_spec=[
            paddle.static.InputSpec(
                shape=[None, None], dtype="int64"),  # input_ids
            paddle.static.InputSpec(
                shape=[None, None], dtype="int64")  # token_type_ids
        ])
    # save inference model
    paddle.jit.save(model, os.path.join(args.save_inference_dir, "inference"))
    tokenizer.save_pretrained(args.save_inference_dir)
    print(
        f"inference model and tokenizer have been saved into {args.save_inference_dir}"
    )


if __name__ == "__main__":
    args = get_args()
    export(args)

[yingyibiao]

LGTM