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.