Minimum and Mean Euclidean distance between two tensors of different shape

username: Tanay_Rastogi
cooked: <p>I am quite new to Pytorch and currently running into issues with Memory Overflow.</p>
<p><strong>Task:</strong><br>
I have two 2D tensors of respective shapes  A: [1000, 14]  &amp;  B: [100000, 14].</p>
<p>I have to find the distance of each row of tensor-A from all rows from tensor-B.  Later using the calculated distance values, I find the mean of minimum/mean distance of each row of tensor-A from tensor-B.</p>
<p><strong>Current Solution:</strong><br>
My solution to calculate minimum distance:</p>
<pre><code class="lang-auto">dist = list()
for row_id in range(A.shape[0]):
      # Mean distance of a row in A from B
      dist.append(torch.linalg.norm(A[row_id, :] - B, dim=1).min().item())
result = torch.FloatTensor(dist).mean()
</code></pre>
<p>And solution to calculate minimum mean distance:</p>
<pre><code class="lang-auto">dist = list()
for row_id in range(A.shape[0]):
      # Mean distance of a row in A from B
      dist.append(torch.linalg.norm(A[row_id, :] - B, dim=1).mean().item())
result = torch.FloatTensor(dist).mean()
</code></pre>
<p><strong>Issue:</strong><br>
This gives me result but is either very slow (if run on CPU) or often leads to memory overflow in GPU when trying to run on GPU. (I have a T4 GPU - 8GB)</p>
<p>Can you please recommend me a better solution to calculate the Euclidean distance that is faster and does not lead to overflow issues?</p>
<p>Thanks!</p>

username: KFrank
cooked: <p>Hi Tanay!</p>
<aside class="quote no-group" data-username="Tanay_Rastogi" data-post="1" data-topic="188598" 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/tanay_rastogi/48/64045_2.png" class="avatar"> Tanay_Rastogi:</div>
<blockquote>
<p>Can you please recommend me a better solution to calculate the Euclidean distance that is faster and does not lead to overflow issues?</p>
</blockquote>
</aside>
<p>Use <code>torch.cdist()</code>.</p>
<p>Consider this script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())

_ = torch.manual_seed (2023)

A = torch.randn (1000, 14, device = 'cuda')
B = torch.randn (100000, 14, device = 'cuda') 
print ('A.shape:         ', A.shape)
print ('B.shape:         ', B.shape)

dist = list()
for row_id in range(A.shape[0]):
    # Min distance of a row in A from B
    dist.append(torch.linalg.norm(A[row_id, :] - B, dim=1).min().item())

resultMinMean = torch.FloatTensor(dist).mean()

dist = list()
for row_id in range(A.shape[0]):
    # Mean distance of a row in A from B
    dist.append(torch.linalg.norm(A[row_id, :] - B, dim=1).mean().item())

resultMeanMean = torch.FloatTensor(dist).mean()

distAB = torch.cdist (A, B)
resultMinMeanB = distAB.min (dim = 1).values.mean()
resultMeanMeanB = distAB.mean()

print ('resultMinMean:   ', resultMinMean)
print ('resultMinMeanB:  ', resultMinMeanB)
print ('allclose? :      ', torch.allclose (resultMinMeanB.cpu(), resultMinMean))

print ('resultMeanMean:  ', resultMeanMean)
print ('resultMeanMeanB: ', resultMeanMeanB)
print ('allclose? :      ', torch.allclose (resultMeanMeanB.cpu(), resultMeanMean))
</code></pre>
<p>And its output:</p>
<pre><code class="lang-plaintext">2.0.1
11.8
GeForce GTX 1050 Ti
A.shape:          torch.Size([1000, 14])
B.shape:          torch.Size([100000, 14])
resultMinMean:    tensor(1.8380)
resultMinMeanB:   tensor(1.8380, device='cuda:0')
allclose? :       True
resultMeanMean:   tensor(5.2006)
resultMeanMeanB:  tensor(5.2006, device='cuda:0')
allclose? :       True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Tanay_Rastogi
cooked: <p>Thanks a lot! This works like a charm. Also, it gave me a very good speed boost as well to my code. Thanks!</p>