PyTorch列表在GPU上切片的速度比在CPU上慢

如何解决PyTorch列表在GPU上切片的速度比在CPU上慢

我想在NVIDIA Jetson Xavier NX（开发套件）上优化ML代码（PyTorch中的SSD）。瓶颈之一似乎是在GPU设备上的PyTorch（1.6.0）张量上进行列表切片。

在NVIDIA GeForce GTX 1050 Ti（GP107）上也发生了同样的问题，cpu速度提高了约2倍。

让我先创建变量

import torch
from time import time

cuda0 = torch.device('cuda:0')

probs = torch.ones([3000],dtype=torch.float64,device=cuda0)
mask = torch.ones([3000],dtype=torch.bool,device=cuda0)

probs_cpu = probs.cpu()
mask_cpu = mask.cpu()

然后运行逻辑（每次运行都会产生大约相同的结果）

before = time()
probs[mask]
print(f'GPU {time() - before:.5f}') # output: GPU 0.00263


before = time()
probs_cpu[mask_cpu]
print(f'cpu {time() - before:.5f}') # output: cpu 0.00066

与上面根据代码在NVIDIA Jetson Xavier NX Developer套件上使用PyTorch库vesrion 1.6.0的cpu相比，为什么列表在GPU上的切片速度要慢4倍？如何加快速度？

代码详细信息：请参阅predictor.py的SSD Implementation in PyTorch中的第51行

在cpu上运行它？：如果我在cpu上运行它，整个算法将不会更快，因为从GPU上下载的时间太长（〜0.00805 s）。