How to calculate pair-wise differences between two tensors in a vectorized way?

username: Yuerno
cooked: <p>I have two tensors of shape (4096, 3) and (4096,3). What I’d like to do is calculate the pairwise differences between all of the individual vectors in those matrices, such that I end up with a (4096, 4096, 3) tensor. This can be done in for-loops, but I’d like to do a vectorized approach. NumPy lets you do some broadcasting approaches, but I’m not sure how to do the same for PyTorch. Any help would be greatly appreciated!</p>

username: ptrblck
cooked: <p>I’m not sure if I understand it correctly, but I think this will do it:</p>
<pre><code class="lang-python">a = torch.randn(5, 3)
b = torch.randn(5, 3)

res = a.unsqueeze(1) - b
# res[i] corresponds to (a[i] - b)
</code></pre>

username: InnovArul
cooked: <p>you can use the code from here for pairwise L2 distance.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/jiyanggao/Video-Person-ReID/blob/master/losses.py#L47-L89" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/jiyanggao/Video-Person-ReID/blob/master/losses.py#L47-L89" target="_blank" rel="nofollow noopener">jiyanggao/Video-Person-ReID/blob/master/losses.py#L47-L89</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="47" style="counter-reset: li-counter 46 ;">
<li>class TripletLoss(nn.Module):</li>
<li>"""Triplet loss with hard positive/negative mining.</li>
<li>
</li>
<li>Reference:</li>
<li>Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.</li>
<li>
</li>
<li>Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.</li>
<li>
</li>
<li>Args:</li>
<li>    margin (float): margin for triplet.</li>
<li>"""</li>
<li>def __init__(self, margin=0.3):</li>
<li>    super(TripletLoss, self).__init__()</li>
<li>    self.margin = margin</li>
<li>    self.ranking_loss = nn.MarginRankingLoss(margin=margin)</li>
<li>
</li>
<li>def forward(self, inputs, targets):</li>
<li>    """</li>
<li>    Args:</li>
<li>        inputs: feature matrix with shape (batch_size, feat_dim)</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/jiyanggao/Video-Person-ReID/blob/master/losses.py#L47-L89" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: Yuerno
cooked: <p>Thanks, this is exactly what I was looking for! Out of curiosity, why is it that the extra dimension is added to the middle, so it becomes (5, 1, 3)?</p>
<p>Lastly, a follow-up question, is it possible to apply this batchwise, so if you had for example:</p>
<pre><code class="lang-auto">a = torch.randn(1, 5, 3)
b = torch.randn(1, 5, 3)
</code></pre>
<p>What seems to be working is just doing <code>res = a.unsqueeze(2) - b</code>, but I’m not sure if that’s 100% correct (although the values match with both methods when I use two fixed arrays).</p>

username: Yuerno
cooked: <p>Thank you for linking the code!</p>

username: ptrblck
cooked: <p>This shouldn’t work for batch sizes &gt; 1.<br>
However, this code should work:</p>
<pre><code class="lang-python">a = torch.randn(2, 5, 3)
b = torch.randn(2, 5, 3)

res = a.unsqueeze(2) - b.unsqueeze(1)
# res[batch_idx, i] corresponds to (a[batch_idx, i] - b[batch_idx])
</code></pre>
<p>The extra dimension is added so that the tensors will automatically be broadcasted.<br>
Basically this would be equal to <code>(a.unsqueeze(1) - b.unsqueeze(0))</code>.</p>

username: Yuerno
cooked: <p>Great, thanks so much!</p>

username: baxincs
cooked: <p>Great answer!<br>
How would I achieve pairwise concatenation?<br>
The ‘res.shape’ should be torch.Size([2, 5, 5, 6]) or torch.Size([2, 10, 3, 3]).</p>

username: ptrblck
cooked: <p>Could you give an example of this operation or post some pseudo code for it, please?</p>

username: baxincs
cooked: <p>Let’s assume that we have a tensor with shape [batch_size, channel_dim, spatial_dim]:</p>
<pre><code class="lang-auto">a = torch.randn(2, 3, 5)
tensor([[[-0.8526,  1.0994, -0.2659, -0.3834,  1.3758],
         [ 0.4018, -0.9505,  0.0167, -0.2450, -1.1768],
         [ 0.9758, -0.7121,  0.7073, -2.0957, -0.8515]],

        [[ 0.4197, -0.0777, -0.0946,  0.1594, -0.0315],
         [-0.1422, -0.8626,  2.9035, -0.5575, -0.4993],
         [-0.7888, -0.5184, -0.7460, -0.4184, -0.2338]]])
</code></pre>
<p>The ‘res’ should have the shape of torch.Size([2, 6, 5, 5]) in this case, and given any pair of spatial locations, the pairwise concatenation result should be:</p>
<pre><code class="lang-auto">res[0, :, 2, 3] = torch.cat([a[0, :, 2], a[0, :, 3]], dim=-1) 
                = tensor([-0.2659,  0.0167,  0.7073, -0.3834, -0.2450, -2.0957])
</code></pre>