Update on "[pytorch] Add triplet margin loss with custom distance"

Summary: As discussed [here](#43342),
adding in a Python-only implementation of the triplet-margin loss that takes a
custom distance function.  Still discussing whether this is necessary to add to
PyTorch Core.

Test Plan: python test/run_tests.py

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: [D23363898](https://our.internmc.facebook.com/intern/diff/D23363898)

[ghstack-poisoned]

test/test_nn.py

@@ -12176,7 +12176,8 @@ def test_triplet_margin_with_distance_loss(self, device):
         def cosine_distance(x, y):
             return 1.0 - F.cosine_similarity(x, y)
 
-        distance_functions = (pairwise_distance, cosine_distance)
+        distance_functions = (pairwise_distance, cosine_distance,
+                              lambda x, y: 1.0 - F.cosine_similarity(x, y))
 
         reductions = ('mean', 'none', 'sum')
         margins = (1.0, 1.5, 0.5)

torch/nn/modules/loss.py

@@ -1337,6 +1337,13 @@ class TripletMarginWithDistanceLoss(_Loss):
     >>>     nn.TripletMarginWithDistanceLoss(distance_function=l_infinity, margin=1.5)
     >>> output = triplet_loss(anchor, positive, negative)
     >>> output.backward()
+    >>>
+    >>> # Custom Distance Function (Lambda)
+    >>> triplet_loss = \
+    >>>     nn.TripletMarginWithDistanceLoss(
+    >>>         distance_function=lambda x, y: 1.0 - F.cosine_similarity(x, y))
+    >>> output = triplet_loss(anchor, positive, negative)
+    >>> output.backward()
 
     Reference:
         V. Balntas, et al.: Learning shallow convolutional feature descriptors with triplet losses: