beta)torch_npu.contrib.module.ROIAlign

接口原型

torch_npu.contrib.module.ROIAlign(nn.Module):

功能描述

使用NPU API进行ROIAlign。

参数说明

输出说明

Float - shape为(k, 5)和(k, 1)的张量。标签以0为基础。

支持的型号

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
>>> from torch_npu.contrib.module import ROIAlign
>>> input1 = torch.FloatTensor([[[[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, 27, 28, 29, 30], [31, 32, 33, 34, 35, 36]]]]).npu()
>>> roi = torch.tensor([[0, -2.0, -2.0, 22.0, 22.0]]).npu()
>>> output_size = (3, 3)
>>> spatial_scale = 0.25
>>> sampling_ratio = 2
>>> aligned = False
>>> input1.requires_grad = True
>>> roi.requires_grad = True
>>> model = ROIAlign(output_size, spatial_scale, sampling_ratio, aligned=aligned).npu()
>>> output = model(input1, roi)
>>> output.sum().backward()