功能描述
- 算子功能:计算高斯误差线性单元的激活函数。
- 计算公式:
Gaussian Error Linear Unit(GELU)的表达式为:

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

接口原型
torch_npu.npu_gelu(Tensor input, approximate='none') -> Tensor
参数说明
- input: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类型,且数据类型必须和input一样,数据格式支持ND,shape必须和input一样, 支持非连续的Tensor。输入最大支持8维。
- Atlas 训练系列产品:数据类型支持FLOAT16、FLOAT32。
- Atlas A2 训练系列产品/Atlas 800I A2 推理产品:数据类型支持FLOAT32、FLOAT16、BFLOAT16。
- Atlas 推理系列产品:数据类型支持FLOAT16、FLOAT32。
约束说明
- 该接口支持图模式(目前仅支持PyTorch 2.1版本)。
- input输入不能含有空指针。
支持的型号
Atlas 训练系列产品
- Atlas A2 训练系列产品/Atlas 800I A2 推理产品
- Atlas 推理系列产品
调用示例
- 单算子模式调用
| import sys
import torch
import torch_npu
input_tensor = torch.randn(100, 200)
output_tensor = torch_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_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")
|