Python接口文档

PAI-Blade提供Python接口供您调用,从而将优化模型集成在工作流中。本文介绍这些Python接口的详细信息,包括签名格式、输入参数及返回结果。

optimize

PAI-Blade主要使用optimize接口进行模型优化,接口详情如下:

  • 签名格式

    def optimize(
        model: Any,
        optimization_level: str,
        device_type: str,
        config: Optional[Config] = None,
        inputs: Optional[List[str]] = None,
        outputs: Optional[List[str]] = None,
        input_shapes: Optional[List[List[str]]] = None,
        input_ranges: Optional[List[List[str]]] = None,
        test_data: List[Dict[str, np.ndarray]] = [],
        calib_data: List[Dict[str, np.ndarray]] = [],
        custom_ops: List[str] = [],
        verbose: bool = False,
    ) -> Tuple[Any, OptimizeSpec, OptimizeReport]:
        pass
  • 输入参数

    参数

    类型

    是否必选

    描述

    默认值

    model

    多种

    待优化的模型,支持以下类型的模型:

    • 如果优化TensorFlow模型,支持以下格式:

      • GraphDef对象。

      • GraphDef PB文件路径,即以.pb.pbtxt为后缀的文件。

      • SavedModel目录,字符串格式的路径。

    • 如果优化PyTorch模型,支持以下格式:

      • torch.nn.Module对象。

      • torch.nn.Module导出的文件,字符串格式的路径,该文件以.pt为后缀。

    optimization_level

    STRING

    优化级别,取值如下(不区分大小写):

    • o1:表示无损优化,包括图改写、编译优化等。

    • o2:表示量化优化。

    device_type

    STRING

    模型运行的设备类型,取值如下(不区分大小写):

    • gpu

    • cpu

    • edge (暂不支持PyTorch)

    inputs

    LIST[STRING]

    输入节点名称。如果不指定该参数,则系统尝试自动推断。

    None

    outputs

    LIST[STRING]

    输出节点名称。如果不指定该参数,则系统尝试自动推断。

    None

    input_shapes

    LIST[LIST[STRING]]

    输入Tensor可能的形状,用于提升特定场景的优化效果。内层列表元素个数必须等于模型输入Tensor的数量,每个元素是一个字符串,表示某种输入的形状,例如'1*512'。如果有多组可能的形状,则在外层列表中增加元素即可。例如,有两个输入的模型可能有如下形状或若干组可能的形状,取值示例如下:

    • [['1*512', '3*256']]

    • [
          ['1*512', '3*256'],
          ['5*512', '9*256'],
          ['10*512', '27*256']
      ]

    None

    input_ranges

    LIST[LIST[STRING]]

    每个输入Tensor的元素取值范围。内层列表元素数量必须等于模型输入Tensor的数量,每个元素是一个字符串,表示某种取值范围。

    取值范围可以使用方括号加实数或字符表示,例如'[1,2]'、'[0.3,0.9]'及'[a,f]'等。如果有多组可能的取值范围,则在外层列表中增加元素即可。例如,有两个输入的模型可能有如下取值范围或若干组可能的取值范围,取值示例分别如下:

    • [['[0.1,0.4]', '[a,f]']]

    • [
          ['[0.1,0.4]', '[a,f]'],
          ['[1.1,1.4]', '[h,l]'],
          ['[2.1,2.4]', '[n,z]']
      ]

    None

    test_data

    多种

    用于模型执行速度校准的测试数据。对于不同类型的模型,其测试数据格式存在差异,其中:

    • TensorFlow模型的测试数据为包含若干组feed_dict的列表,类型为LIST[DICT[STRING, np.ndarray]]。

    • PyTorch模型的测试数据为若干组输入Tensor Tuple,类型为LIST[Tuple[torch.tensor, ]]。

    []

    calib_data

    多种

    用于量化优化的校准数据。当优化级别为o2时,必须指定该参数。校准数据的数据类型与测试数据相同。

    []

    custom_ops

    LIST[STRING]

    自定义算子库的路径。如果模型中包含用户的自定义算子库文件,需要把依赖的库文件路径放到列表中。

    []

    verbose

    BOOL

    是否打印更多日志,包括如下取值:

    • True:打印更多日志

    • False:无需打印更多日志

    False

    config

    blade.Config

    高级配置,详情请参见下方的blade.Config参数表。

    blade.Config用于传递优化的高级参数,其构造函数的签名如下。

    class Config(ABC):
        def __init__(
            self,
            disable_fp16_accuracy_check: bool = False,
            disable_fp16_perf_check: bool = False,
            enable_static_shape_compilation_opt: bool = False,
            enable_dynamic_shape_compilation_opt: bool = True,
            quant_config: Optional[Dict[str, str]] = None,
        ) -> None:
            pass

    签名中的参数含义如下表所示。

    表 1. blade.Config

    参数

    是否必选

    类型

    描述

    默认值

    disable_fp16_accuracy_check

    BOOL

    是否在fp16优化中开启准确性校验,取值如下:

    • False:关闭准确性校验

    • True:开启准确性校验

    False

    disable_fp16_perf_check

    BOOL

    是否在fp16优化中开启性能校验,取值如下:

    • False:关闭性能校验

    • True:开启性能校验

    False

    enable_static_shape_compilation_opt

    BOOL

    是否开启固定尺寸的编译优化,即Static Shape模式编译,取值如下:

    • False:关闭固定尺寸的编译优化

    • True:开启固定尺寸的编译优化

    False

    enable_dynamic_shape_compilation_opt

    BOOL

    是否开启动态尺寸的编译优化,即Dynamic Shape模式编译,取值如下:

    • False:关闭动态尺寸的编译优化

    • True:开启动态尺寸的编译优化

    True

    quant_config

    DICT[STRING, STRING]

    量化优化相关的配置。目前仅支持唯一的keyweight_adjustment,表示是否尝试调整模型参数以减小量化精度损失。该key值对应的value包括如下取值:

    • "true":开启该选项。

    • "false":关闭该选项。

    None

  • 返回值

    返回一个具有三个元素的Tuple,类型为Tuple[Any, OptimizeSpec, OptimizeReport]。其中三个元素的详细信息如下表所示。

    序号

    含义

    类型

    描述

    优化后的模型

    多种

    其类型与输入模型相同。例如TensorFlow SavedModel的返回结果是GraphDef对象。

    优化的外部依赖

    OptimizeSpec

    Python中通过with语句使达到预期优化效果的外部依赖都生效,该依赖包括环境变量、编译缓存等。SDK中暂不需要该参数。

    优化报告

    OptimizeReport

    详情请参见优化报告