add triplet margin loss (pytorch#1165)

Author: edgarriba

test/test_nn.py

@@ -1828,6 +1828,20 @@ def test_pairwise_distance(self):
         input2 = Variable(torch.randn(4, 4), requires_grad=True)
         self.assertTrue(gradcheck(lambda x, y: F.pairwise_distance(x, y), (input1, input2)))
 
+    def test_triplet_margin_loss(self):
+        input1 = Variable(torch.randn(4, 4), requires_grad=True)
+        input2 = Variable(torch.randn(4, 4), requires_grad=True)
+        input3 = Variable(torch.randn(4, 4), requires_grad=True)
+        self.assertTrue(gradcheck(lambda x1, x2, x3: F.triplet_margin_loss(
+            x1, x2, x3), (input1, input2, input3)))
+
+    def test_triplet_margin_swap_loss(self):
+        input1 = Variable(torch.randn(4, 4), requires_grad=True)
+        input2 = Variable(torch.randn(4, 4), requires_grad=True)
+        input3 = Variable(torch.randn(4, 4), requires_grad=True)
+        self.assertTrue(gradcheck(lambda x1, x2, x3: F.triplet_margin_loss(
+            x1, x2, x3, swap=True), (input1, input2, input3)))
+
 
 class TestNNInit(TestCase):
     def setUp(self):

torch/nn/functional.py

@@ -634,8 +634,9 @@ def pairwise_distance(x1, x2, p=2, eps=1e-6):
             \Vert x \Vert _p := \left( \sum_{i=1}^n  \vert x_i \vert ^ p \right) ^ {1/p}
 
         Args:
-            x (Tensor): input tensor containing the two input batches
-            p (real): the norm degree. Default: 2
+            x1: first input tensor
+            x2: second input tensor
+            p: the norm degree. Default: 2
 
         Shape:
             - Input: :math:`(N, D)` where `D = vector dimension`
@@ -651,3 +652,55 @@ def pairwise_distance(x1, x2, p=2, eps=1e-6):
     diff = torch.abs(x1 - x2)
     out = torch.pow(diff + eps, p).sum(dim=1)
     return torch.pow(out, 1. / p)
+
+
+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.
+    This is used for measuring a relative similarity between samples. A triplet is composed by
+    `a`, `p` and `n`: anchor, positive examples and negative example respectively.
+    The shape of all input variables should be :math:`(N, D)`.
+
+    The distance swap is described in detail in the paper `Learning shallow convolutional feature descriptors with
+    triplet losses`_ by V. Balntas, E. Riba et al.
+
+    .. math::
+        L(a, p, n) = \frac{1}{N} \left( \sum_{i=1}^N \max \{d(a_i, p_i) - d(a_i, n_i) + {\rm margin}, 0\} \right)
+
+    where :math: `d(x_i, y_i) = \| {\bf x}_i - {\bf y}_i \|_2^2`.
+
+    Args:
+        anchor: anchor input tensor
+        positive: positive input tensor
+        negative: negative input tensor
+        p: the norm degree. Default: 2
+        eps: small epsilon value to avoid numerical issues
+        swap: compute distance swap
+
+    Shape:
+        - Input: :math:`(N, D)` where `D = vector dimension`
+        - Output: :math:`(N, 1)`
+
+        >>> input1 = autograd.Variable(torch.randn(100, 128))
+        >>> input2 = autograd.Variable(torch.randn(100, 128))
+        >>> input3 = autograd.Variable(torch.randn(100, 128))
+        >>> output = F.triplet_margin_loss(input1, input2, input3, p=2)
+        >>> output.backward()
+
+    .. _Learning shallow convolutional feature descriptors with triplet losses:
+        http://www.iis.ee.ic.ac.uk/%7Evbalnt/shallow_descr/TFeat_paper.pdf
+    """
+    assert anchor.size() == positive.size(), "Input sizes between positive and negative must be equal."
+    assert anchor.size() == negative.size(), "Input sizes between anchor and negative must be equal."
+    assert positive.size() == negative.size(), "Input sizes between positive and negative must be equal."
+    assert anchor.dim() == 2, "Inputd must be a 2D matrix."
+    assert margin > 0.0, 'Margin should be positive value.'
+    d_p = pairwise_distance(anchor, positive, p, eps)
+    d_n = pairwise_distance(anchor, negative, p, eps)
+    if swap:
+        d_s = pairwise_distance(positive, negative, p, eps)
+        d_n = torch.min(d_n, d_s)
+
+    dist_hinge = torch.clamp(margin + d_p - d_n, min=0.0)
+    loss = torch.mean(dist_hinge)
+    return loss

torch/nn/modules/loss.py

@@ -427,6 +427,53 @@ def forward(self, input, target):
                                              self.margin, weight=self.weight)(input, target)
 
 
+class TripletMarginLoss(Module):
+    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.
+    This is used for measuring a relative similarity between samples. A triplet is composed by
+    `a`, `p` and `n`: anchor, positive examples and negative example respectively.
+    The shape of all input variables should be :math:`(N, D)`.
+
+    The distance swap is described in detail in the paper `Learning shallow convolutional feature descriptors with
+    triplet losses`_ by V. Balntas, E. Riba et al.
+
+    .. math::
+        L(a, p, n) = \frac{1}{N} \left( \sum_{i=1}^N \max \{d(a_i, p_i) - d(a_i, n_i) + {\rm margin}, 0\} \right)
+
+    where :math: `d(x_i, y_i) = \| {\bf x}_i - {\bf y}_i \|_2^2`.
+
+    Args:
+        anchor: anchor input tensor
+        positive: positive input tensor
+        negative: negative input tensor
+        p: the norm degree. Default: 2
+
+    Shape:
+        - Input: :math:`(N, D)` where `D = vector dimension`
+        - Output: :math:`(N, 1)`
+
+    >>> triplet_loss = nn.TripletMarginLoss(margin=1.0, p=2)
+    >>> input1 = autograd.Variable(torch.randn(100, 128))
+    >>> input2 = autograd.Variable(torch.randn(100, 128))
+    >>> input3 = autograd.Variable(torch.randn(100, 128))
+    >>> output = triplet_loss(input1, input2, input3)
+    >>> output.backward()
+
+    .. _Learning shallow convolutional feature descriptors with triplet losses:
+        http://www.iis.ee.ic.ac.uk/%7Evbalnt/shallow_descr/TFeat_paper.pdf
+    """
+
+    def __init__(self, margin=1.0, p=2, eps=1e-6, swap=False):
+        super(TripletMarginLoss, self).__init__()
+        self.margin = margin
+        self.p = p
+        self.eps = eps
+        self.swap = swap
+
+    def forward(self, anchor, positive, negative):
+        return F.triplet_margin_loss(anchor, positive, negative, self.margin,
+                                     self.p, self.eps, self.swap)
+
 # TODO: L1HingeEmbeddingCriterion
 # TODO: MSECriterion weight
 # TODO: ClassSimplexCriterion