文档HOT
下载
中文
注册

torch_npu.npu_gelu

功能描述

  • 算子功能：计算高斯误差线性单元的激活函数。
  • 计算公式：

    Gaussian Error Linear Unit(GELU)的表达式为：

    Φ(x)是Gaussian Distribution的CDF(Cumulative Distribution Function), 表达式为：

接口原型

1
torch_npu.npu_gelu(Tensor self, approximate='none') -> Tensor

参数说明

  • self：Device侧的Tensor类型，待进行npu_gelu计算的入参，数据格式支持ND，支持非连续的Tensor。输入最大支持8维。
    • Atlas 训练系列产品：数据类型支持FLOAT16、FLOAT32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT32、FLOAT16、BFLOAT16。
    • Atlas 推理系列产品：数据类型支持FLOAT16、FLOAT32。
  • approximate：可选参数，计算使用的激活函数模式，字符串类型，可配置为"none"或者"tanh"，其中"none"代表使用erf模式，"tanh"代表使用tanh模式。

输出说明

  • y： Device侧的Tensor类型，且数据类型必须和self一样，数据格式支持ND，shape必须和self一样， 支持非连续的Tensor。输入最大支持8维。
    • Atlas 训练系列产品：数据类型支持FLOAT16、FLOAT32。
    • Atlas A2 训练系列产品/Atlas 800I A2 推理产品：数据类型支持FLOAT32、FLOAT16、BFLOAT16。
    • Atlas 推理系列产品：数据类型支持FLOAT16、FLOAT32。

约束说明

  • 该接口支持图模式（目前仅支持PyTorch 2.1版本）。
  • self输入不能含有空指针。

支持的型号

  • Atlas 训练系列产品

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

调用示例

  • 单算子模式调用
    1
2
3
4
5
6
    import sys
import torch
import torch_npu

input_tensor = torch.randn(100, 200)
output_tensor = torch.ops.npu.npu_gelu(input_tensor.npu(), approximate='tanh')
  • 图模式调用
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
    # 入图方式
import os
import torch
import torch_npu
import numpy as np
import torch.nn as nn
import torchair as tng
from torchair.configs.compiler_config import CompilerConfig

torch_npu.npu.set_compile_mode(jit_compile=True)
class Net(torch.nn.Module):

    def __init__(self):
        super().__init__()
    def forward(self, self_0, approximate):
        out = torch.ops.npu.npu_gelu(self_0, approximate=approximate)
        return out


x = torch.randn(100, 10, 20).npu()

model = Net()
config = CompilerConfig()
npu_backend = tng.get_npu_backend(compiler_config=config)
model = torch.compile(model, fullgraph=True, backend=npu_backend, dynamic=False)
npu_out = model(x, approximate="none")
父主题： torch_npu