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）。

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）