由以下代码进行实验,发现:
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)