昇腾社区首页
中文
注册

torch_npu.npu_anti_quant

功能描述

  • 算子功能:将INT8数据反量化为FP16。
  • 计算公式为:

接口原型

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

参数说明

  • x:Tensor类型,即输入参数中的x。数据类型支持INT8,数据格式支持ND,支持非连续的Tensor。输入最大支持8维。
  • scale:Tensor类型,数据类型支持FLOAT32、BFLOAT16,数据格式支持ND,支持非连续的Tensor,仅支持1维Tensor。
  • offset:Tensor类型,可选参数,数据类型支持FLOAT32、BFLOAT16,数据格式支持ND,支持非连续的Tensor,仅支持1维Tensor,且shape必须与scale的shape大小一致。
  • dst_dtype:ScalarType类型,可选参数,默认值为torch.float16。
  • src_dtype:ScalarType类型,可选参数,输入值默认为torch.int8。

输出说明

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

约束说明

  • x、scale这两个输入中不能含有空指针。
  • 如果输入scale的shape值不为1,则输入x的最后一维shape值必须与scale的shape一致。
  • 数据类型BFLOAT16仅如下产品型号支持:

    Atlas A2训练系列产品/Atlas 800I A2推理产品

支持的PyTorch版本

  • PyTorch 2.2
  • PyTorch 2.1

支持的型号

  • Atlas A2训练系列产品/Atlas 800I A2推理产品

调用示例

  • 单算子调用
    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)
  • 图模式调用(图模式目前仅支持PyTorch 2.1版本)
    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)