Add F.cosine_similarity (pytorch#1502)

Author: apaszke

docs/source/nn.rst

@@ -160,7 +160,7 @@ Pooling Layers
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: AdaptiveMaxPool2d
-    :members:       
+    :members:
 
 :hidden:`AdaptiveAvgPool1d`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -174,7 +174,7 @@ Pooling Layers
 .. autoclass:: AdaptiveAvgPool2d
     :members:
 
-       
+
 Non-linear Activations
 ----------------------------------
 
@@ -682,7 +682,7 @@ Pooling functions
 
 .. autofunction:: adaptive_avg_pool2d
 
-   
+
 Non-linear activation functions
 -------------------------------
 
@@ -814,6 +814,11 @@ Distance functions
 
 .. autofunction:: pairwise_distance
 
+:hidden:`cosine_similarity`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: cosine_similarity
+
 
 Loss functions
 --------------

test/test_nn.py

@@ -2063,6 +2063,16 @@ def test_triplet_margin_swap_loss(self):
         self.assertTrue(gradcheck(lambda x1, x2, x3: F.triplet_margin_loss(
             x1, x2, x3, swap=True), (input1, input2, input3)))
 
+    def test_cosine_similarity(self):
+        input1 = Variable(torch.randn(4, 4), requires_grad=True)
+        input2 = Variable(torch.randn(4, 4), requires_grad=True)
+        self.assertTrue(gradcheck(lambda x, y: F.cosine_similarity(x, y), (input1, input2)))
+
+        input1 = Variable(torch.randn(4, 5, 6), requires_grad=True)
+        input2 = Variable(torch.randn(4, 5, 6), requires_grad=True)
+        self.assertTrue(gradcheck(lambda x, y: F.cosine_similarity(x, y, dim=0), (input1, input2)))
+        self.assertTrue(gradcheck(lambda x, y: F.cosine_similarity(x, y, dim=-1), (input1, input2)))
+
     def test_bilinear(self):
         module = nn.Bilinear(10, 10, 8)
         module2 = legacy.Bilinear(10, 10, 8)

torch/nn/functional.py

@@ -694,6 +694,25 @@ def pairwise_distance(x1, x2, p=2, eps=1e-6):
     return torch.pow(out, 1. / p)
 
 
+def cosine_similarity(x1, x2, dim=1, eps=1e-8):
+    r"""Returns cosine similarity between x1 and x2, computed along dim.
+
+    Args:
+        x1 (Variable): First input.
+        x2 (Variable): Second input (of size matching x1).
+        dim (int, optional): Dimension of vectors. Default: 1
+        eps (float, optional): Small value to avoid division by zero. Default: 1e-8
+
+    Shape:
+        - Input: :math:`(\ast_1, D, \ast_2)` where D is at position `dim`.
+        - Output: :math:`(\ast_1, \ast_2)` where 1 is at position `dim`.
+    """
+    w12 = torch.sum(x1 * x2, dim)
+    w1 = torch.norm(x1, 2, dim)
+    w2 = torch.norm(x2, 2, dim)
+    return (w12 / (w1 * w2).clamp(min=eps)).squeeze()
+
+
 def triplet_margin_loss(anchor, positive, negative, margin=1.0, p=2, eps=1e-6, swap=False):
     r"""Creates a criterion that measures the triplet loss given an input tensors x1, x2, x3
     and a margin with a value greater than 0.