Apply a function along an axis

username: thomas
cooked: <p>Hello everyone!</p>
<p>I have an input of this shape:</p>
<pre><code class="lang-auto">(num_samples, num_symbols, num_features, num_observations)
</code></pre>
<p>I would like to feed through this input through a neural network.</p>
<p>The neural network as <code>(num_features, num_observations)</code> shape input and <code>(num_outputs)</code> outputs, giving me <code>(num_samples, num_symbols, num_outputs)</code> when I apply along</p>
<pre><code class="lang-auto">def apply_along_axis(function, x, axis: int = 0):
    return torch.stack([
        function(x_i) for x_i in torch.unbind(x, dim=axis)
    ], dim=axis)
</code></pre>
<p>(from <a href="https://discuss.pytorch.org/t/pytorch-equivalent-of-numpy-apply-along-axis/3422" class="inline-onebox">Pytorch equivalent of numpy.apply_along_axis</a> )</p>
<p>This is an operation that is made very often in my network, and I just discover that 42% of the time is spent in the apply_along_axis.</p>
<p>I wanted to know if there is better way to do this ? (the for loop seems to be quite slow)</p>
<p>Is there a possibility to create identical versions of the same “function” in the network and update all the weights automatically when grad is invoked ?</p>
<p>(I maybe have a similar question with more complex situation if there is good suggestion related to this issue)</p>
<p>Thanks!</p>

username: KFrank
cooked: <p>Hi Thomas!</p>
<aside class="quote no-group quote-modified" data-username="thomas" data-post="1" data-topic="130440" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e36b37/48.png" class="avatar"> thomas:</div>
<blockquote>
<p>I have an input of this shape:</p>
<pre data-code-wrap="jsx"><code class="lang-jsx">(num_samples, num_symbols, num_features, num_observations)
</code></pre>
<p>I would like to feed through this input through a neural network.</p>
<p>…</p>
<pre data-code-wrap="jsx"><code class="lang-jsx">def apply_along_axis(function, x, axis: int = 0):
    return torch.stack([
        function(x_i) for x_i in torch.unbind(x, dim=axis)
    ], dim=axis)
</code></pre>
<p>I wanted to know if there is better way to do this ?</p>
</blockquote>
</aside>
<p>I do not know of any functionality built into pytorch similar to your<br>
<code>apply_along_axis()</code>.  And even if there were, it would still impose<br>
a performance penalty, as it would still be breaking up what might<br>
have been a single, larger tensor operation into many smaller, axis-wise<br>
tensor operations (absent some kind of hypothetical JIT compilation).</p>
<p>As a general rule, if you find yourself looping over a tensor, you should<br>
see if you can recast your computation into pure (that is, loop-free)<br>
tensor operations.  Sometimes you can and sometimes you can’t.</p>
<p>Note that you can <em>sometimes</em> realize a net performance gain by getting<br>
rid of loops, <em>even if</em> your loop-free approach is of higher computational<br>
complexity.</p>
<p>Here is a simplistic, contrived example that illustrates replacing<br>
<code>apply_along_axis()</code> with a single pytorch tensor operation:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt; def apply_along_axis(function, x, axis: int = 0):
...     return torch.stack([
...         function(x_i) for x_i in torch.unbind(x, dim=axis)
...     ], dim=axis)
...
&gt;&gt;&gt; def my_fn (x):
...     return torch.softmax (x, 0)
...
&gt;&gt;&gt; t = torch.randn (2, 3)
&gt;&gt;&gt; t
tensor([[ 2.2871,  0.6413, -0.8615],
        [-0.3649, -0.6931,  0.9023]])
&gt;&gt;&gt; apply_along_axis (my_fn, t, 0)
tensor([[0.8092, 0.1561, 0.0347],
        [0.1897, 0.1366, 0.6737]])
&gt;&gt;&gt; torch.softmax (t, 1)
tensor([[0.8092, 0.1561, 0.0347],
        [0.1897, 0.1366, 0.6737]])
&gt;&gt;&gt; apply_along_axis (my_fn, t, 1)
tensor([[0.9341, 0.7916, 0.1463],
        [0.0659, 0.2084, 0.8537]])
&gt;&gt;&gt; torch.softmax (t, 0)
tensor([[0.9341, 0.7916, 0.1463],
        [0.0659, 0.2084, 0.8537]])
</code></pre>
<p>You suggest that your use case involves having the <code>function</code> you apply<br>
be an entire neural-network model.</p>
<p>Although various model layers do have constraints on the shapes they<br>
expect, the basic building blocks accept (and sometimes require) an<br>
arbitrary batch dimension.  This suggests that reworking your model so<br>
that you don’t need <code>apply_along_axis</code> could be plausible.</p>
<p>Two building-block examples:  <code>Linear</code> accepts an arbitrary number<br>
of leading “batch” dimensions, so that’s likely to be easy.  On the other<br>
hand, <code>Conv2d</code> requires a tensor of exactly four dimensions, but it’s<br>
leading dimension is an arbitrary batch dimension, so you can use<br>
<code>view()</code> (or <code>reshape()</code>) to repackaged multiple “batch” dimensions<br>
into a single batch dimension.  Thus:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; torch.nn.Linear (2, 3) (torch.randn (7, 7, 5, 2)).shape
torch.Size([7, 7, 5, 3])
&gt;&gt;&gt; torch.nn.Conv2d (2, 3, 3, padding = 1) (torch.randn (7, 7, 5, 2, 11, 11).view (7*7*5, 2, 11, 11)).view (7, 7, 5, 3, 11, 11).shape
torch.Size([7, 7, 5, 3, 11, 11])
</code></pre>
<p>Whether or not you would be able to push these kinds of techniques<br>
through an entire model will depend on the model’s details, but there<br>
are certainly some realistic models where you could.</p>
<p>Best.</p>
<p>K. Frank</p>