Caffe2 - Python API A deep learning, cross platform ML framework
distance.py
1 import torch
2 from .module import Module
3 from .. import functional as F
4 from ..._jit_internal import weak_module, weak_script_method
5
6
7 @weak_module
8 class PairwiseDistance(Module):
9  r"""
10  Computes the batchwise pairwise distance between vectors :math:v_1, :math:v_2 using the p-norm:
11
12  .. math ::
13  \Vert x \Vert _p = \left( \sum_{i=1}^n \vert x_i \vert ^ p \right) ^ {1/p}
14
15  Args:
16  p (real): the norm degree. Default: 2
17  eps (float, optional): Small value to avoid division by zero.
18  Default: 1e-6
19  keepdim (bool, optional): Determines whether or not to keep the batch dimension.
20  Default: False
21
22  Shape:
23  - Input1: :math:(N, D) where D = vector dimension
24  - Input2: :math:(N, D), same shape as the Input1
25  - Output: :math:(N). If :attr:keepdim is False, then :math:(N, 1).
26
27  Examples::
28
29  >>> pdist = nn.PairwiseDistance(p=2)
30  >>> input1 = torch.randn(100, 128)
31  >>> input2 = torch.randn(100, 128)
32  >>> output = pdist(input1, input2)
33  """
34  __constants__ = ['norm', 'eps', 'keepdim']
35
36  def __init__(self, p=2., eps=1e-6, keepdim=False):
37  super(PairwiseDistance, self).__init__()
38  self.norm = p
39  self.eps = eps
40  self.keepdim = keepdim
41
42  @weak_script_method
43  def forward(self, x1, x2):
44  return F.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim)
45
46
47 @weak_module
48 class CosineSimilarity(Module):
49  r"""Returns cosine similarity between :math:x_1 and :math:x_2, computed along dim.
50
51  .. math ::
52  \text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)}
53
54  Args:
55  dim (int, optional): Dimension where cosine similarity is computed. Default: 1
56  eps (float, optional): Small value to avoid division by zero.
57  Default: 1e-8
58
59  Shape:
60  - Input1: :math:(\ast_1, D, \ast_2) where D is at position dim
61  - Input2: :math:(\ast_1, D, \ast_2), same shape as the Input1
62  - Output: :math:(\ast_1, \ast_2)
63
64  Examples::
65
66  >>> input1 = torch.randn(100, 128)
67  >>> input2 = torch.randn(100, 128)
68  >>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)
69  >>> output = cos(input1, input2)
70  """
71  __constants__ = ['dim', 'eps']
72
73  def __init__(self, dim=1, eps=1e-8):
74  super(CosineSimilarity, self).__init__()
75  self.dim = dim
76  self.eps = eps
77
78  @weak_script_method
79  def forward(self, x1, x2):
80  return F.cosine_similarity(x1, x2, self.dim, self.eps)