gather与grid_sample对比

由以下代码进行实验，发现：

gather运行时间为：
0.7160608768463135S

grid_sample bilinear

0.7564301490783691S

grid_sample nearest

0.5S

import torch
import time
import torch.nn.functional as F
s=torch.rand(9,32,160,160)
s=s.view(9,32,160*160)
p=torch.rand(9,32,160,160)
z=torch.ones(9,32,160,160).float()
z=z.view(9,32,160*160)
r=torch.ones(9,32,160,160).long()
r=r.view(9,32,160*160)
out_w=160
out_h=160
y,x=torch.meshgrid(torch.arange(0,out_h),torch.arange(out_w))
y=y-(out_h-1)/2
y=y/((out_h-1)/2)

x=x-(out_w-1)/2
x=x/((out_w-1)/2)
grid = torch.cat((x.unsqueeze(2), y.unsqueeze(2)), dim=2)
grid = grid.unsqueeze(0).expand(9,-1,-1,-1)
st=time.time()
for i in range(40):
    # g=torch.gather(s,dim=2,index=r)
    g = F.grid_sample(p, grid=grid, mode='bilinear',align_corners=True,padding_mode='reflection')
ed=time.time()
print(ed-st)


st=time.time()
for i in range(40):
     g=torch.gather(s,dim=2,index=r)
    #q = F.grid_sample(p, grid=grid, mode='nearest',align_corners=True,padding_mode='reflection')
ed=time.time()
print(ed-st)