Fast interp for CPU and GPU
import torch
+from copy import deepcopy
from operator import mul
from torch.autograd import Variable
from .interpolation import Interpolation
@@ -90,6 +91,8 @@ def left_interp(interp_indices, interp_values, rhs):
return res
else:
+ # Special cuda version -- this is faster on the GPU for some reason
+ if interp_indices.is_cuda:
if interp_indices.ndimension() == 3:
n_batch, n_data, n_interp = interp_indices.size()
interp_indices = interp_indices.view(-1)
@@ -125,6 +128,16 @@ def left_interp(interp_indices, interp_values, rhs):
res = res.sum(-2)
return res
+ # Special non-cuda version -- this is faster on the CPU
+ else:
+ interp_size = list(interp_indices.size()) + [rhs.size(-1)]
+ rhs_size = deepcopy(interp_size)
+ rhs_size[-3] = rhs.size()[-2]
+ interp_indices_expanded = interp_indices.unsqueeze(-1).expand(*interp_size)
+ res = rhs.unsqueeze(-2).expand(*rhs_size).gather(-3, interp_indices_expanded)
+ res = res.mul(interp_values.unsqueeze(-1).expand(interp_size))
+ return res.sum(-2)
+
def sparse_eye(size):
"""
