API说明

PyTorch-NPU 除了提供了 PyTorch 官方算子实现之外，也提供了大量高性能的自定义算子，详细的算子信息以及描述如下所示：

备注

在运行下述示例之前，需要导入torch_npu扩展包 import torch_npu

torch_npu._npu_dropout(self, p)

不使用种子(seed)进行dropout结果计数,与torch.dropout相似，优化NPU设备实现

参数:

  • self (Tensor) -- 输入张量

  • p (Float) -- 丢弃概率

返回类型:

(Tensor, Tensor)

示例:

 1>>> input = torch.tensor([1.,2.,3.,4.]).npu()
 2>>> input
 3tensor([1., 2., 3., 4.], device='npu:0')
 4>>> prob = 0.3
 5>>> output, mask = torch_npu._npu_dropout(input, prob)
 6>>> output
 7tensor([0.0000, 2.8571, 0.0000, 0.0000], device='npu:0')
 8>>> mask
 9tensor([ 98, 255, 188, 186, 120, 157, 175, 159,  77, 223, 127,  79, 247, 151,
10    253, 255], device='npu:0', dtype=torch.uint8)
torch_npu.copy_memory_(dst, src, non_blocking=False) → Tensor

从src拷贝元素到self张量，并返回self

约束说明:

copy_memory_仅支持NPU张量,copy_memory_的输入张量应具有相同的dtype和设备index

参数:

  • dst (Tensor) -- 拷贝源张量

  • sr (Tensor) -- 返回张量所需数据类型

  • non_blocking (Bool,Default: False) -- 如果设置为True且此拷贝位于CPU和NPU之间，则拷贝可能相对于主机异步发生,在其他情况下，此参数没有效果

返回类型:

Tensor

示例:

1>>> a=torch.IntTensor([0,  0, -1]).npu()
2>>> b=torch.IntTensor([1, 1, 1]).npu()
3>>> a.copy_memory_(b)
4tensor([1, 1, 1], device='npu:0', dtype=torch.int32)
torch_npu.empty_with_format(size, dtype, layout, device, pin_memory, acl_format)

返回一个填充未初始化数据的张量

参数:

  • size (ListInt) -- 定义输出张量shape的整数序列,可以是参数数量(可变值)，也可以是列表或元组等集合

  • dtype (torch.dtype,Default: None) -- 返回张量所需数据类型;如果值为None，请使用全局默认值(请参见torch.set_default_tensor_type()).

  • layout (torch.layout, Default: torch.strided) -- 返回张量所需布局

  • device (torch.device, Default: None) -- 返回张量的所需设备

  • pin_memory (Bool, Default: False) -- 返回张量的所需设备

  • acl_format (Int, Default: 2) -- 返回张量所需内存格式

返回类型:

Tensor

示例:

1>>> torch_npu.empty_with_format((2, 3), dtype=torch.float32, device="npu")
2tensor([[1., 1., 1.],
3        [1., 1., 1.]], device='npu:0')
torch_npu.fast_gelu(self) → Tensor

gelu的npu实现,支持FakeTensor模式

参数:

self (Tensor) -- 输入张量(只float16、float32)

返回类型:

Tensor

示例:

 1# Normal
 2>>> x = torch.rand(2).npu()
 3>>> x
 4tensor([0.5991, 0.4094], device='npu:0')
 5>>> torch_npu.fast_gelu(x)
 6tensor([0.4403, 0.2733], device='npu:0')
 7
 8# FakeTensorMode
 9>>> from torch._subclasses.fake_tensor import FakeTensorMode
10>>> with FakeTensorMode():
11...     x = torch.rand(2).npu()
12...     torch_npu.fast_gelu(x)
13>>> FakeTensor(..., device='npu:0', size=(2,))
torch_npu.npu_alloc_float_status(self) → Tensor

生成一个包含8个0的一维张量

参数:

self (Tensor) -- 输入张量

返回类型:

Tensor

示例:

1>>> input    = torch.randn([1,2,3]).npu()
2>>> output = torch_npu.npu_alloc_float_status(input)
3>>> input
4tensor([[[ 2.2324,  0.2478, -0.1056],
5        [ 1.1273, -0.2573,  1.0558]]], device='npu:0')
6>>> output
7tensor([0., 0., 0., 0., 0., 0., 0., 0.], device='npu:0')
torch_npu.npu_anchor_response_flags(self, featmap_size, stride, num_base_anchors) → Tensor

在单个特征图中生成锚点的责任标志

参数:

  • self (Tensor) -- 真值框，shape为[batch, 4]的2D张量

  • featmap_size (ListInt[2]) -- 特征图大小

  • strides (ListInt[2]) -- 当前水平的步长

  • num_base_anchors (Int) -- base anchors的数量

返回类型:

Tensor

示例:

1>>> x = torch.rand(100, 4).npu()
2>>> y = torch_npu.npu_anchor_response_flags(x, [60, 60], [2, 2], 9)
3>>> y.shape
4torch.Size([32400])