昇腾社区首页
EN
注册

torch_npu.npu_anti_quant

功能描述

  • 算子功能:对张量x进行反量化操作,即将整数恢复为浮点数。
  • 计算公式为:

接口原型

1
torch_npu.npu_anti_quant(Tensor x, Tensor scale, *, Tensor? offset=None, ScalarType? dst_dtype=None, ScalarType? src_dtype=None) -> Tensor

参数说明

  • x:Tensor类型,即输入参数中的x。数据格式支持ND,支持非连续的Tensor。输入最大支持8维。
    • Atlas 推理系列产品:数据类型支持int8。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件:数据类型支持int8、int32,其中int32类型数据的每个值是由8个int4数值拼成。
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品:数据类型支持int8、int32,其中int32类型数据的每个值是由8个int4数值拼成。
  • scale:Tensor类型,反量化参数scale。仅支持1维Tensor,shape为(n,)。其中n可以为1,如果n不为1,当x为int8类型,必须与输入x的尾轴维度大小相同;当x为int32类型时,必须为输入x的尾轴维度大小的8倍。数据格式支持ND,支持非连续的Tensor。
    • Atlas 推理系列产品:数据类型支持float32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件:数据类型支持float32、bfloat16。
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品:数据类型支持float32、bfloat16。
  • offset:Tensor类型,可选参数,反量化参数offset。仅支持1维Tensor,数据类型和shape必须与scale一致。数据格式支持ND,支持非连续的Tensor。
  • dst_dtype:ScalarType类型,可选参数,指定输出的数据类型,默认值为float16。
    • Atlas 推理系列产品:数据类型支持float16。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件:数据类型支持float16、bfloat16。
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品:数据类型支持float16、bfloat16。
  • src_dtype:ScalarType类型,可选参数,指定源输入的数据类型,默认值为int8。
    • Atlas 推理系列产品:数据类型支持int8。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件:数据类型支持quint4x2或int8。
    • Atlas A3 训练系列产品/Atlas A3 推理系列产品:数据类型支持quint4x2或int8。

输出说明

一个Tensor类型的输出,代表antiquant的计算结果。

约束说明

  • 该接口支持图模式(目前仅支持PyTorch 2.1版本)。
  • x、scale这两个输入中不能有空(None)。

支持的型号

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

调用示例

  • 单算子模式调用
    1
    2
    3
    4
    5
    6
    import torch
    import torch_npu
    x_tensor = torch.tensor([1,2,3,4], dtype=torch.int8).npu()
    scale = torch.tensor([2.0], dtype=torch.float).npu()
    offset = torch.tensor([2.0], dtype=torch.float).npu()
    out = torch_npu.npu_anti_quant(x_tensor, scale, offset=offset, dst_dtype=torch.float16)
    
  • 图模式调用
     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    import torch
    import torch_npu
    import torchair as tng
    from torchair.ge_concrete_graph import ge_apis as ge
    from torchair.configs.compiler_config import CompilerConfig
    
    config = CompilerConfig()
    config.debug.graph_dump.type = 'pbtxt'
    npu_backend = tng.get_npu_backend(compiler_config=config)
    x_tensor = torch.tensor([1,2,3,4], dtype=torch.int8).npu()
    scale = torch.tensor([2.0], dtype=torch.float).npu()
    offset = torch.tensor([2.0], dtype=torch.float).npu()
    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
        def forward(self,x,scale,offset):
            return torch_npu.npu_anti_quant(x, scale, offset=offset, dst_dtype=torch.float16)
    cpu_model = Model()
    model = torch.compile(cpu_model, backend=npu_backend, dynamic=False, fullgraph=True)
    output = model(x_tensor,scale,offset)