torch_npu.npu_fast_gelu
功能描述
- 算子功能:快速高斯误差线性单元激活函数(Fast Gaussian Error Linear Units activation function),对输入的每个元素计算FastGelu;输入是具有任何有效形状的张量。
接口原型
torch_npu.npu_fast_gelu(Tensor self) -> Tensor
参数说明
self:Tensor类型,即输入参数中的x。数据类型支持FLOAT16、FLOAT32、BFLOAT16,数据格式支持ND,支持非连续的Tensor。输入最大支持8维。
输出说明
一个Tensor类型的输出,代表fast_gelu的计算结果。
支持的型号
- Atlas A2训练系列产品/Atlas 800I A2推理产品
调用示例
- 单算子调用
1 2 3 4 5 6 7
import os import torch import torch_npu import numpy as np data_var = np.random.uniform(0, 1, [4, 2048, 16, 128]).astype(np.float32) x = torch.from_numpy(data_var).to(torch.float32).npu() y = torch_npu.npu_fast_gelu(x).cpu().numpy()
- 图模式调用(图模式目前仅支持PyTorch 2.1版本)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
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 os.environ["ENABLE_ACLNN"] = "false" torch_npu.npu.set_compile_mode(jit_compile=True) class Network(nn.Module): def __init__(self): super(Network, self).__init__() def forward(self, x): y = torch_npu.npu_fast_gelu(x) return y npu_mode = Network() config = CompilerConfig() npu_backend = tng.get_npu_backend(compiler_config=config) npu_mode = torch.compile(npu_mode, fullgraph=True, backend=npu_backend, dynamic=False) data_var = np.random.uniform(0, 1, [4, 2048, 16, 128]).astype(np.float32) x = torch.from_numpy(data_var).to(torch.float32) y =npu_mode(x).cpu().numpy()
父主题: torch_npu