指定尺寸时，torch.max 使用 GPU 比使用 CPU 慢

Question

t1_h = torch.tensor(np.arange(100000), dtype=torch.float32)
cuda0 = torch.device('cuda:0')
t1_d = torch.tensor(np.arange(100000), dtype=torch.float32, device = cuda0)

%timeit -n 10000 max_h = torch.max(t1_h, 0)
%timeit -n 10000 max_d = torch.max(t1_d, 0)

10000 个循环，最好的 3 个：每个循环 144 µs

10000 个循环，最好的 3 个：每个循环 985 µs

正如你在上面看到的，GPU 比 CPU 花费更多的时间。 但是如果我没有指定计算最大值的维度，那么 GPU 会更快。

%timeit -n 10000 max_h = torch.max(t1_h)
%timeit -n 10000 max_d = torch.max(t1_d)

10000 个循环，最好的 3 个：每个循环 111 µs

10000 个循环，最好的 3 个：每个循环 41.8 µs

我也尝试过使用argmax而不是max但它工作正常（GPU 比 CPU 快）。

%timeit -n 10000 cs_h = torch.argmax(t1_h, 0)
%timeit -n 10000 cs_d = torch.argmax(t1_d, 0)

10000 个循环，最好的 3 个：每个循环 108 µs

10000 个循环，最好的 3 个：每个循环 18.1 µs

指定尺寸后， torch.max在 GPU 上torch.max缓慢有什么原因吗？

Answer 1

我自己发现了这一点，并在 PyTorch 中打开了一个问题。 看起来很快就会修复 - 也许是 1.5 或 1.6 版？ - 但与此同时，建议的解决方法是使用

ii=a.argmax(0)
maxval = a.gather(0, ii.unsqueeze(0)).squeeze(0)