torch_npu.npu_trans_quant_param

功能描述

完成量化计算参数scale数据类型的转换。

接口原型

npu_trans_quant_param(Tensor scale, Tensor? offset=None) -> Tensor

参数说明

输出说明

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

约束说明

  1. 传入的scale或out不能为空指针。
  2. scale、offset或out的数据类型和数据格式需要在支持的范围之内。
  3. offset、scale的shape需要为1维(t,),t = 1或n。

支持的PyTorch版本

支持的型号

Atlas A2 训练系列产品

调用示例

单算子调用:
import torch
import torch_npu
import logging
import os

scale = torch.randn(16, dtype=torch.float32)
offset = torch.randn(16, dtype=torch.float32)
npu_out = torch_npu.npu_quant_matmul(scale.npu(), offset.npu())

图模式:
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
import logging
from torchair.core.utils import logger
logger.setLevel(logging.DEBUG)
import os
import numpy as np
os.environ["ENABLE_ACLNN"] = "true"
config = CompilerConfig()
model = cpu_model.npu()
npu_backend = tng.get_npu_backend(compiler_config=config)


class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self,scale, offset):
        return torch_npu.npu_trans_quant_param(scale, offset)
    
cpu_model = MyModel()
scale = torch.randn(1, dtype=torch.float32)
offset = torch.randn(1, dtype=torch.float32)
model = torch.compile(cpu_model, backend=npu_backend, dynamic=True)
npu_out = model(scale.npu(), offset.npu())