pytorch sort / медиана на GPU так же медленно, как numpy на CPU. Я что-то пропустил?

Question

Мне нужно вычислить медианный массив из списка массивов. Вот временные результаты короткого упражнения. У меня есть графический процессор NVIDIA Quadro P2000, и я проверил с помощью БПФ и сверток, что cuda / torch работает в 100 раз быстрее, чем numpy. Тем не менее, когда дело доходит до сортировки, графический процессор довольно «застрял» и не дает мне значительного ускорения.

import numpy as np
import torch as th
c=np.random.rand(3,480,700).astype(np.float32) # cpu version
g=th.tensor(c).to(device='cuda:0')             # gpu version
%timeit np.median(c,axis=0)        # 14.3 ms ± 61.6 µs per loop 
%timeit np.sort(c,axis=0)[1]       # 9.59 ms ± 26.5 µs per loop
%timeit th.median(g,dim=0).values  # 11.8 ms ± 9.13 µs per loop
%timeit th.sort(g,dim=0).values[1] # 9.28 ms ± 6.01 µs per loop

Я проверил с помощью nvidia-smi, что GPU работает во время вызовов факела, fwiw.

Существует множество отчетов (например, this ), в которых показан алгоритм сортировки факела, отбивающий штаны numpy, но я этого не получаю. Кто-нибудь знает, есть ли простой трюк, чтобы заставить сортировку / медиану факела работать на GPU быстрее?

Answer 1

Что такое x?

Может быть, вы его неправильно разместили? Потому что имеет смысл, если вместо этого c.

Кроме того, как заявило @JoshVarty, реализация pyTorch превосходит numpy, когда вы изменяете ее на x.

Итак, ваш код должен быть:

import numpy as np
import torch as th
c=np.random.rand(3,480,700).astype(np.float32) # cpu version
g=th.tensor(c).to(device='cuda:0')             # gpu version
%timeit np.median(c,axis=0)        
%timeit np.sort(c,axis=0)[1]      
%timeit th.median(g,dim=0).values 
%timeit th.sort(g,dim=0).values[1]