How to calculate distances between classes

Podidiving: <p>Hello, everyone</p>
<p>I have the next code:</p>
<pre><code class="lang-auto">feature = torch.randn(10, 20)
target = torch.randint(0, 4, (10,))
total = 0
for i in range(0, 4):
  for j in range(i + 1, 4):
    diff = feature[target == i].mean(0) - feature[target == j].mean(0)
    total += torch.norm(diff, 2)

</code></pre>
<p>Can you help me to optimise it? I feel, it can be done without <code>for</code> loops, but I don’t understand how to get rid of them</p>
<p>Thanks in advance</p>

KFrank: <p>Hi Podi!</p>
<aside class="quote no-group" data-username="Podidiving" data-post="1" data-topic="149524" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/podidiving/48/31083_2.png" class="avatar"> Podidiving:</div>
<blockquote>
<p>I feel, it can be done without <code>for</code> loops</p>
</blockquote>
</aside>
<p>You need to address two issues to avoid explicit loops:</p>
<p>First, you have to rework <code>feature[target == i]</code>.  This is because the<br>
<code>shape</code> of <code>feature[target == i]</code> depends on the value of <code>i</code>.  Because<br>
pytorch does not support <em>ragged tensors</em> (tensors whose slices don’t<br>
have the same shape), you won’t be able to replace this version of the<br>
loop with operations on a single tensor.</p>
<p>Therefore we will mask <code>feature</code> rather than indexing into it.</p>
<p>Second your <code>i</code> and <code>j</code> loops index over the upper triangle of the matrix<br>
you would get if you wrote <code>diff</code> out in matrix form.  Pytorch does not<br>
(directly) support triangular matrices.</p>
<p>Therefore we will compute the full square <code>diff</code> matrix (which will be<br>
symmetric, so there will be some redundancy, but that’s okay), and use<br>
<code>triu()</code> to give us just the elements we want to sum over.</p>
<p>Thus:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.10.2
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; feature = torch.randn (10, 20)
&gt;&gt;&gt; target = torch.randint (0, 4, (10,))
&gt;&gt;&gt;
&gt;&gt;&gt; total = 0
&gt;&gt;&gt; for i in range(0, 4):
...   for j in range(i + 1, 4):
...     diff = feature[target == i].mean(0) - feature[target == j].mean(0)
...     total += torch.norm(diff, 2)
...
&gt;&gt;&gt; total
tensor(32.2318)
&gt;&gt;&gt;
&gt;&gt;&gt; mask = target.unsqueeze (-1) == torch.arange (4)
&gt;&gt;&gt; means = ((feature.unsqueeze (1) * mask.unsqueeze (-1)).sum (0) / mask.sum (0).unsqueeze (-1))
&gt;&gt;&gt; totalB = torch.norm (means.unsqueeze (0) - means.unsqueeze (1), 2, dim = 2).triu().sum()
&gt;&gt;&gt; totalB
tensor(32.2318)
</code></pre>
<p>As an aside, if <code>target == i</code> sums to zero for any value of <code>i</code>, <code>.mean()</code> won’t<br>
really be defined, and you will get <code>nan</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; target
tensor([3, 0, 0, 0, 2, 3, 3, 1, 0, 1])
&gt;&gt;&gt; target[4] = 1
&gt;&gt;&gt; (target == 2).sum()
tensor(0)
&gt;&gt;&gt; total = 0
&gt;&gt;&gt; for i in range(0, 4):
...   for j in range(i + 1, 4):
...     diff = feature[target == i].mean(0) - feature[target == j].mean(0)
...     total += torch.norm(diff, 2)
...
&gt;&gt;&gt; total
tensor(nan)
&gt;&gt;&gt; mask = target.unsqueeze (-1) == torch.arange (4)
&gt;&gt;&gt; means = ((feature.unsqueeze (1) * mask.unsqueeze (-1)).sum (0) / mask.sum (0).unsqueeze (-1))
&gt;&gt;&gt; totalB = torch.norm (means.unsqueeze (0) - means.unsqueeze (1), 2, dim = 2).triu().sum()
&gt;&gt;&gt; totalB
tensor(nan)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>