昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

torch_npu.npu_quantize

功能说明

  • 算子功能：对输入的张量进行量化处理。
  • 计算公式：

    • 如果div_mode为True：

    • 如果div_mode为False：

函数原型

[object Object]

参数说明

  • input (Tensor)：需要进行量化的源数据张量，必选输入，数据格式支持NDND，支持非连续的Tensor。div_mode为False且dtype为quint4x2时，最后一维需要能被8整除。

    • [object Object]Atlas 推理系列产品[object Object]：数据类型支持float、float16。
    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：数据类型支持float、float16、bfloat16。

  • scales (Tensor)：对input进行scales的张量，必选输入：

    • div_mode 为True时

      • [object Object]Atlas 推理系列产品[object Object]：数据类型支持float。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：数据类型支持float、bfloat16。

    • div_mode 为False时，数据格式支持NDND，支持非连续的Tensor。支持1维或多维（1维时，对应轴的大小需要与input中第axis维相等或等于1；多维时，scales的shape需要与input的shape维度相等，除axis指定的维度，其他维度为1，axis指定的维度必须和input对应的维度相等或等于1）。

      • [object Object]Atlas 推理系列产品[object Object]：数据类型支持float、float16。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：数据类型支持float、float16、bfloat16。

  • zero_points (Tensor)：对input进行offset的张量，可选输入。

    • div_mode 为True时

      • [object Object]Atlas 推理系列产品[object Object]：数据类型支持int8、uint8、int32。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：数据类型支持int8、uint8、int32、bfloat16。

    • div_mode 为False时，数据格式支持NDND，支持非连续的Tensor。支持1维或多维（1维时，对应轴的大小需要与input中第axis维相等或等于1；多维时，scales的shape需要与input维度相等，除axis指定的维度，其他维度为1，axis指定的维度必须和input对应的维度相等）。zero_points的shape和dtype需要和scales一致。

      • [object Object]Atlas 推理系列产品[object Object]：数据类型支持float、float16。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：数据类型支持float、float16、bfloat16。

  • dtype (int)：int类型，指定输出参数的类型。

    • div_mode 为True时，

      • [object Object]Atlas 推理系列产品[object Object]：类型支持qint8、quint8、int32。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：类型支持qint8、quint8、int32。

    • div_mode 为False时，类型支持qint8、quint4x2。如果dtype为quint4x2时，输出tensor类型为int32，由8个int4拼接。

  • axis (int)：int类型，量化的element-wise轴，其他的轴做broadcast，默认值为1。

    div_mode 为False时，axis取值范围是[-2, +∞）且指定的轴不能超过输入input的维度数。如果axis=-2，代表量化的element-wise轴是输入input的倒数第二根轴；如果axis大于-2，量化的element-wise轴是输入的最后一根轴。

  • div_mode (bool)：布尔类型，表示计算scales模式。当div_mode为True时，表示用除法计算scales；div_mode为False时，表示用乘法计算scales，默认值为True。

返回值

Tensor

公式中的输出，输出大小与input一致。数据类型由参数dtype指定，如果参数dtype为quint4x2，输出的dtype是int32，shape的最后一维是输入shape最后一维的1/8，shape其他维度和输入一致。

约束说明

  • 该接口支持推理场景下使用。
  • 该接口支持图模式（PyTorch 2.1版本）。
  • div_mode为False时：
    • 支持[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]。
    • 当dtype为quint4x2或者axis为-2时，不支持[object Object]Atlas 推理系列产品[object Object]。

支持的型号

  • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]
  • [object Object]Atlas 推理系列产品[object Object]

调用示例

  • 单算子模式调用

    • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]

      [object Object]

    • [object Object]Atlas 推理系列产品[object Object]

      [object Object]

  • 图模式调用

    [object Object]