Speed up torch.narrow over multiple dimensions

username: weiskohlmoe
cooked: <p>I want to implement a Machine Learning model that, among other things, should learn coefficient for (Quadratic) spline functions. (<a href="https://en.wikipedia.org/wiki/B-spline" class="inline-onebox" rel="noopener nofollow ugc">B-spline - Wikipedia</a> for reference). As a part of the forward operation the spline function is applied on the input.</p>
<p>For the computation value of this function I implemented the algorithm as described here:</p><aside class="onebox wikipedia" data-onebox-src="https://en.wikipedia.org/wiki/De_Boor%27s_algorithm">
  <header class="source">

      <a href="https://en.wikipedia.org/wiki/De_Boor%27s_algorithm" target="_blank" rel="noopener nofollow ugc">en.wikipedia.org</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://en.wikipedia.org/wiki/De_Boor%27s_algorithm" target="_blank" rel="noopener nofollow ugc">De Boor's algorithm</a></h3>

<p>In the mathematical subfield of numerical analysis de Boor's algorithm is a polynomial-time and numerically stable algorithm for evaluating spline curves in B-spline form. It is a generalization of de Casteljau's algorithm for Bézier curves. The algorithm was devised by Carl R. de Boor.  Simplified, potentially faster variants of the de Boor algorithm have been created but they suffer from comparatively lower stability.
 A general introduction to B-splines is given in the main article. Here we di...</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
In the initialization step I need to get values out of a tensor of coefficients according to a tensor of indices.<br>
Since I need a range of Indices I used <code>torch.narrow</code> for this problem.<br>
However my runtime is dominated by this computation, so I am searching for ways to speed it up</p>
<p>My current code to measure the runtime looks like this:</p>
<pre><code class="lang-auto">import torch
import time
import torch.nn.functional as F

num_weights = 31
batch_size = 24
values_size = 45 * 45
coeffs = torch.rand(num_weights)
values = torch.rand(batch_size, values_size)
interval = torch.arange(0, 1.1, 0.1)

#This is just so the snippet below works without error
index_tensor = torch.searchsorted(interval, values, right=True, side="right") + 2*torch.ones(values.size(), dtype=torch.int64)
pd = (2,2)
interval = interval.unsqueeze(0)
interval = F.pad(interval, pd, "replicate")
interval = interval.squeeze()


start = time.time()
coeff_matrix = torch.stack([torch.stack([torch.narrow(coeffs, 0, index_tensor[j][i] - 2, 3) for i in range(values_size)]) for j in range(batch_size)])
print("Time to build coeff_matrix", time.time() - start)
</code></pre>
<p>On my computer I get a runtime of ~0.5 seconds. Just for reference, the other operations in the evaluation amount to ~0.01 seconds<br>
I tried to calculate the matrix first without <code>torch.stack()</code> in the middle, but it was slower.<br>
So my question is, if there is a faster method to compute my <code>coeff_matrix</code>  I thought I could use<code>torch.gather(),</code>but I am unsure about how to do it.<br>
Additionally I thought I could abuse the fact that the coefficients I want to take depend on i and j linearly, but I don’t know how.<br>
The coefficients are also parameters in the real application, so it is important to not break the gradient.</p>

username: weiskohlmoe
cooked: <p>My Idea about using <code>torch.gather()</code> worked out fine for me. Here is the solution I used:</p>
<pre><code class="lang-auto">bigger_coeffs = coeffs.unsqueeze(0).unsqueeze(0)

bigger_coeffs = coeffs.expand(batch_size, values_size, -1)

bigger_index_tensor = index_tensor.unsqueeze(-1).expand(index_tensor.shape[0], index_tensor.shape[1], 3)

sub = torch.tensor([2, 1, 0])

index_tensor_res = bigger_index_tensor - sub[None, None, :]

coeff_matrix = torch.gather(bigger_coeffs, 2, index_tensor_res)
</code></pre>