torchvision::nms不支持在npu侧计算

问题：
执行算子torchvision::nms时，会切换到cpu侧运行，运行速度慢，会出现警告
[W compiler_depend.ts:55] Warning: CAUTION: The operator 'torchvision::nms' is not currently supported on the NPU backend and will fall back to run on the CPU. This may have performance implications. (function npu_cpu_fallback)

环境：
pytorch 2.1.0
cann 8.0.rc1.alpha002
npu-smi version 23.0.rc2.2
python 3.9
torch                         2.1.0
torch-npu                     2.1.0+gitb88f102
torchair                      0.1
torchvision                   0.16.0
torchvision_npu               0.16.0+gitunknown
euler 2.8.3

运行截图：
![CPU运行日志](https://foruda.gitee.com/images/1721721882620811795/29157e79_14603773.png "2024_07_23_16_04_07_firefox.png")
![NPU运行日志](https://foruda.gitee.com/images/1721721916539982890/950292f4_14603773.png "2024_07_23_16_05_01_firefox.png")

测试代码：
import torch
import torch_npu
import torchvision
import torchvision_npu
import time

device = torch.device("npu:0")
boxes = torch.rand(1000, 4).to(device) * 1000
scores = torch.rand(1000).to(device)
iou_threshold = 0.5

start_time = time.perf_counter()
keep = torchvision.ops.nms(boxes, scores, iou_threshold)
torch_npu.npu.synchronize()
end_time = time.perf_counter()
run_time = end_time - start_time
print(f"Execution time: {run_time:.6f} seconds")

Ascend/vision
暂停

内容风险标识

评论 (4)

Ascend/vision暂停 .gitee-modal { width: 500px !important; }

内容风险标识