How to compute the cosine normalization distance using nn.CosineSimilarity?

John1231983: <p>I have two tensors <code>x</code> and <code>y</code>. I want to compute normalized cosine distance <code>d_norm</code> as follows<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/e/1e04cf96a23114d1388d5388ea7250a43769ee4d.png" data-download-href="https://discuss.pytorch.org/uploads/default/1e04cf96a23114d1388d5388ea7250a43769ee4d" title="Screenshot%20from%202019-10-17%2016-46-27"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/e/1e04cf96a23114d1388d5388ea7250a43769ee4d_2_517x58.png" alt="Screenshot%20from%202019-10-17%2016-46-27" data-base62-sha1="4hyD96UedwKsfd3btJQU2kBWqCN" width="517" height="58" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/e/1e04cf96a23114d1388d5388ea7250a43769ee4d_2_517x58.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/e/1e04cf96a23114d1388d5388ea7250a43769ee4d_2_775x87.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/e/1e04cf96a23114d1388d5388ea7250a43769ee4d_2_1034x116.png 2x" data-dominant-color="FAFAFA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot%20from%202019-10-17%2016-46-27</span><span class="informations">1483×169 11.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
where<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/a/4a082dbe6c67e0da00c72ff4ed2b78f5082fe34e.png" data-download-href="https://discuss.pytorch.org/uploads/default/4a082dbe6c67e0da00c72ff4ed2b78f5082fe34e" title="Screenshot%20from%202019-10-17%2016-42-40"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a082dbe6c67e0da00c72ff4ed2b78f5082fe34e_2_517x53.png" alt="Screenshot%20from%202019-10-17%2016-42-40" data-base62-sha1="ayUR2O3TdJ3kTYAHRrByrq0tguO" width="517" height="53" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a082dbe6c67e0da00c72ff4ed2b78f5082fe34e_2_517x53.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a082dbe6c67e0da00c72ff4ed2b78f5082fe34e_2_775x79.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/a/4a082dbe6c67e0da00c72ff4ed2b78f5082fe34e_2_1034x106.png 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot%20from%202019-10-17%2016-42-40</span><span class="informations">1219×126 20.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Assume that the term <code>x_i - mu_y</code> and <code>y_j - mu_y</code>replaced by <code>x_normalized</code> and <code>y_normalized</code>, respectively. The <code>d_norm</code>  can be computed as follows</p>
<pre><code class="lang-auto">N, C,H,W = x_normalized.size()
x_normalized = x_normalized.reshape(N, C, -1)                                # (N, C, H*W)
y_normalized = y_normalized.reshape(N, C, -1)                                # (N, C, H*W)
cosine_sim = torch.bmm(x_normalized.transpose(1, 2), y_normalized)           # (N, H*W, H*W)
d = 1 - cosine_sim                                  # (N, H*W, H*W)  d[n, i, j] means d_ij for n-th data 
d_min, _ = torch.min(d, dim=2, keepdim=True)        # (N, H*W, 1)
d_norm = d / (d_min + 1e-5)
</code></pre>
<p>The above code works but it has one issue is that CUDA memory error when the size of <code>x_normalized</code> and <code>y_normalized</code> are big such as 256x256. That reason is due to the <code>torch.bmm</code> function. I am wonder do we have another way for batch multiplication as bmm did? If not, can we use the built-in function <a href="https://pytorch.org/docs/master/nn.html#torch.nn.CosineSimilarity" rel="noopener nofollow ugc">torch.nn.CosineSimilarity</a> to compute the d_norm?<br>
Thanks.</p>
<p>This is what I did using the pytorch built in function. However, two results are different.</p>
<pre><code class="lang-auto">cos sim: tensor(-6.5344e-05, grad_fn=&lt;MeanBackward0&gt;)
tensor(36834.2500, grad_fn=&lt;MeanBackward0&gt;)

Using pytorch function
cos sim: tensor(-0.0559, grad_fn=&lt;MeanBackward0&gt;)
tensor(0.9977, grad_fn=&lt;MeanBackward0&gt;)
</code></pre>
<p>To reproduce the result, you can use the code</p>
<pre><code class="lang-auto">import torch
import numpy as np
import torch.nn as nn

def cosine_dist(x, y, h=0.5):
        assert x.size() == y.size()
        N, C, H, W = x.size() 
        y_mu = y.mean(3).mean(2).mean(0).reshape(1, -1, 1, 1)
        x_centered = x - y_mu
        y_centered = y - y_mu
        x_normalized = x_centered / torch.norm(x_centered, p=2, dim=1, keepdim=True)
        y_normalized = y_centered / torch.norm(y_centered, p=2, dim=1, keepdim=True)

        # The equation at the bottom of page 6 in the paper
        # Vectorized computation of cosine similarity for each pair of x_i and y_j
        x_normalized = x_normalized.reshape(N, C, -1)                                # (N, C, D*H*W)
        y_normalized = y_normalized.reshape(N, C, -1)                                # (N, C, D*H*W)
        cosine_sim = torch.bmm(x_normalized.transpose(1, 2), y_normalized)           # (N, D*H*W, D*H*W)
        print ('cos sim:', torch.mean(cosine_sim))
        d = 1 - cosine_sim                                  # (N, D*H*W, D*H*W)  d[n, i, j] means d_ij for n-th data 
        d_min, _ = torch.min(d, dim=2, keepdim=True)        # (N, D*H*W, 1)

        d_norm = d / (d_min + 1e-5)       
        return d_norm

def cosine_dist_pytorch(x, y, h=0.5):
        assert x.size() == y.size()
        N, C, H, W = x.size()   
        y_mu = y.mean(3).mean(2).mean(0).reshape(1, -1, 1, 1)

        x_centered = x - y_mu
        y_centered = y - y_mu
        x_normalized = x_centered / torch.norm(x_centered, p=2, dim=1, keepdim=True)
        y_normalized = y_centered / torch.norm(y_centered, p=2, dim=1, keepdim=True)
     
        x_normalized = x_normalized.reshape(N, C, -1)                                # (N, C, D*H*W)
        y_normalized = y_normalized.reshape(N, C, -1)                                # (N, C, D*H*W)
        #cosine_sim = torch.bmm(x_normalized.transpose(1, 2), y_normalized)           # (N, D*H*W, D*H*W)
        cos = nn.CosineSimilarity(dim=1, eps=1e-6)
        cosine_sim = cos(x_normalized, y_normalized).unsqueeze_(2)
        print ('cos sim:', torch.mean(cosine_sim))
        d = 1 - cosine_sim                                  # (N, D*H*W, D*H*W)  d[n, i, j] means d_ij for n-th data 
        d_min, _ = torch.min(d, dim=2, keepdim=True)        # (N, D*H*W, 1)
        d_norm = d / (d_min + 1e-5)       
        return d_norm 
      
img1 = torch.rand((4, 2, 16, 16), requires_grad=True)
img2 = torch.rand((4, 2, 16, 16), requires_grad=True)

loss1 = cosine_dist(img1, img2)
print (torch.mean(loss1))
print ('\nUsing pytorch function')
loss2 = cosine_dist_pytorch(img1, img2)
print (torch.mean(loss2))
</code></pre>