Slow loop. better method?

Mark_Ettinger: <p>I’ve defined a custom loss function for a regression problem but it contains a loop that is very slow and I was wondering if there is a better way to do this.</p>
<p>I generate a prediction and then I want to find the vector in a table (i.e. embedding layer) which is closest to the residual.  I use a loop for this but this is very slow.  I wanted to use torch.Tensor.apply_ but this is for cpu only, not gpu.  I discovered torch.func.vmap but I’m not sure how to use this over both the table and a minibatch (which seems to be its primary purpose).  I should also mention I want to update the table during training so perhaps I need to use requires_grad = True in torch.arange?</p>
<p>Thank you in advance for any advice!</p>
<pre><code class="lang-auto">def lossFunction(self, prediction, label):
        residual = (label - prediction)
        for i in torch.arange(self.nCentroids, device='cuda'):
            thisNorm = torch.linalg.vector_norm(residual - self.embedding(i))
            if i == 0:
                bestNorm = thisNorm
                bestIndex = 0
            elif thisNorm &lt; bestNorm:
                bestNorm = thisNorm
                bestIndex = i
        thisNorm = torch.linalg.vector_norm(prediction - self.embedding(i))
        return thisNorm
</code></pre>

KFrank: <p>Hi Mark!</p>
<aside class="quote no-group" data-username="Mark_Ettinger" data-post="1" data-topic="219030">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mark_ettinger/48/76602_2.png" class="avatar"> Mark_Ettinger:</div>
<blockquote>
<p>I want to find the vector in a table (i.e. embedding layer) which is closest to the residual. I use a loop for this but this is very slow.</p>
</blockquote>
</aside>
<p>If you can package your embedding-layer table as a pytorch tensor, you can find the<br>
nearest vector with pytorch tensor operations, avoiding an explicit python loop.</p>
<aside class="quote no-group" data-username="Mark_Ettinger" data-post="1" data-topic="219030">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mark_ettinger/48/76602_2.png" class="avatar"> Mark_Ettinger:</div>
<blockquote>
<pre><code class="lang-auto">            thisNorm = torch.linalg.vector_norm(residual - self.embedding(i))
</code></pre>
</blockquote>
</aside>
<p>Here, <code>self.embedding</code> looks like a function that takes an index (packaged as a<br>
zero-dimensional pytorch tensor).  Let me assume that <code>embedding</code> is or can be packaged<br>
as a tensor.  In particular, let <code>prediction</code> and <code>residual</code> be vectors of length <code>n</code> and<br>
embedding be a tensor of shape <code>[n, nCentroids]</code>.  You can then use a single call to<br>
<code>argmin()</code> instead of the explicit loop.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        thisNorm = torch.linalg.vector_norm(prediction - self.embedding(i))
</code></pre>
</blockquote>
</aside>
<p>This appears to be a typo – I assume that this should be:</p>
<pre data-code-wrap="python"><code class="lang-python">        thisNorm = torch.linalg.vector_norm (prediction - self.embedding (bestIndex))
</code></pre>
<p>Note, we will use <em>broadcasting</em> to compare the <code>residual</code> vector with each of the columns<br>
of the <code>embedding</code> tensor.  That is, <code>(residual - embedding).shape)</code> is <code>[n, nCentroids]</code>.</p>
<p>Here is a loop-free script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

torch.manual_seed (2025)

n = 3
nCentroids = 5

prediction = torch.randn (n)
residual = torch.randn (n)
embedding = torch.randn (nCentroids, n)

bestIndex = torch.linalg.vector_norm (residual - embedding, dim = 1).argmin()
thisNorm = torch.linalg.vector_norm (prediction - embedding[bestIndex])

print ('bestIndex:', bestIndex)
print ('thisNorm:', thisNorm)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.6.0+cu126
bestIndex: tensor(3)
thisNorm: tensor(1.4950)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Mark_Ettinger: <p>Thank you KFrank!  I have two questions.</p>
<p>In order to make the table trainable, it must be nn.Parameter rather than torch.Tensor, right?</p>
<p>So I created self.table as nn.Parameter but in this line:</p>
<pre><code class="lang-auto">bestIndex = torch.linalg.vector_norm (residual - self.table, dim = 1).argmin()
</code></pre>
<p>I get this error (batch size is 32 and the table is 65536 x 19):</p>
<p>RuntimeError: The size of tensor a (32) must match the size of tensor b (65536) at non-singleton dimension 0</p>
<p>I could fix the batch size and hardcode the broadcast but I think perhaps I’m missing something important?</p>

KFrank: <p>Hi Mark!</p>
<aside class="quote no-group" data-username="Mark_Ettinger" data-post="3" data-topic="219030">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mark_ettinger/48/76602_2.png" class="avatar"> Mark_Ettinger:</div>
<blockquote>
<p>In order to make the table trainable, it must be nn.Parameter rather than torch.Tensor, right?</p>
</blockquote>
</aside>
<p>This would be the typical way to do it.  (There would be a connotation that <code>self.embedding</code><br>
is naturally part of your model, but this is not a requirement.)  But you can also train a tensor<br>
that carries <code>requires_grad = True</code> without it being a <code>Parameter</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>I get this error (batch size is 32 and the table is 65536 x 19):</p>
<p>RuntimeError: The size of tensor a (32) must match the size of tensor b (65536) at non-singleton dimension 0</p>
</blockquote>
</aside>
<p>Take a look at the documentation for <a href="https://pytorch.org/docs/2.6/generated/torch.linalg.vector_norm.html" rel="noopener nofollow ugc">linalg.vector_norm()</a> and <a href="https://pytorch.org/docs/2.6/notes/broadcasting.html" rel="noopener nofollow ugc">broadcasting semantics,</a><br>
in particular how the <code>dim</code> argument works for <code>vector_norm()</code>.  Print out the <code>shape</code>s for<br>
<code>residual</code> and <code>self.table</code> and see if everything lines up correctly.</p>
<p>Best.</p>
<p>K. Frank</p>