mirror of
https://github.com/ml-explore/mlx.git
synced 2025-07-26 04:21:17 +08:00
Triplet Loss (#211)
* Triplet Loss * Requested Changes * Margin to alpha
This commit is contained in:
parent
dfa9f4bc58
commit
e549f84532
@ -16,3 +16,4 @@ Loss Functions
|
|||||||
mse_loss
|
mse_loss
|
||||||
nll_loss
|
nll_loss
|
||||||
smooth_l1_loss
|
smooth_l1_loss
|
||||||
|
triplet_loss
|
@ -232,6 +232,48 @@ def smooth_l1_loss(
|
|||||||
return _reduce(loss, reduction)
|
return _reduce(loss, reduction)
|
||||||
|
|
||||||
|
|
||||||
|
def triplet_loss(
|
||||||
|
anchors: mx.array,
|
||||||
|
positives: mx.array,
|
||||||
|
negatives: mx.array,
|
||||||
|
axis: int = -1,
|
||||||
|
p: int = 2,
|
||||||
|
margin: float = 1.0,
|
||||||
|
eps: float = 1e-6,
|
||||||
|
reduction: str = "none",
|
||||||
|
) -> mx.array:
|
||||||
|
r"""
|
||||||
|
Computes the triplet loss for a set of anchor, positive, and negative samples.
|
||||||
|
Margin is represented with alpha in the math section.
|
||||||
|
|
||||||
|
.. math::
|
||||||
|
|
||||||
|
L_{\text{triplet}} = \max\left(\|A - P\|_p - \|A - N\|_p + \alpha, 0\right)
|
||||||
|
|
||||||
|
Args:
|
||||||
|
anchors (array): The anchor samples.
|
||||||
|
positives (array): The positive samples.
|
||||||
|
negatives (array): The negative samples.
|
||||||
|
axis (int, optional): The distribution axis. Default: ``-1``.
|
||||||
|
p (int, optional): The norm degree for pairwise distance. Default: ``2``.
|
||||||
|
margin (float, optional): Margin for the triplet loss. Defaults to ``1.0``.
|
||||||
|
eps (float, optional): Small positive constant to prevent numerical instability. Defaults to ``1e-6``.
|
||||||
|
reduction (str, optional): Specifies the reduction to apply to the output:
|
||||||
|
``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input;
|
||||||
|
if reduction is "mean" or "sum", returns a scalar tensor.
|
||||||
|
"""
|
||||||
|
loss = mx.maximum(
|
||||||
|
mx.sqrt(mx.power(anchors - positives, p).sum(axis) + eps)
|
||||||
|
- mx.sqrt(mx.power(anchors - negatives, p).sum(axis) + eps)
|
||||||
|
+ margin,
|
||||||
|
0,
|
||||||
|
)
|
||||||
|
return _reduce(loss, reduction)
|
||||||
|
|
||||||
|
|
||||||
def _reduce(loss: mx.array, reduction: str = "none"):
|
def _reduce(loss: mx.array, reduction: str = "none"):
|
||||||
if reduction == "mean":
|
if reduction == "mean":
|
||||||
return mx.mean(loss)
|
return mx.mean(loss)
|
||||||
|
@ -239,6 +239,32 @@ class TestNN(mlx_tests.MLXTestCase):
|
|||||||
expected_sum = mx.sum(expected_none)
|
expected_sum = mx.sum(expected_none)
|
||||||
self.assertTrue(mx.allclose(losses_sum, expected_sum))
|
self.assertTrue(mx.allclose(losses_sum, expected_sum))
|
||||||
|
|
||||||
|
def test_triplet_loss(self):
|
||||||
|
anchors = mx.array([[1, 2, 3], [1, 2, 3]])
|
||||||
|
positives = mx.array([[4, 5, 6], [0, -1, 2]])
|
||||||
|
negatives = mx.array([[7, 8, 9], [3, 2, 3]])
|
||||||
|
|
||||||
|
# Test with reduction 'none'
|
||||||
|
losses_none = nn.losses.triplet_loss(
|
||||||
|
anchors, positives, negatives, reduction="none"
|
||||||
|
)
|
||||||
|
expected_none = mx.array([0, 2.31662])
|
||||||
|
self.assertTrue(mx.allclose(losses_none, expected_none))
|
||||||
|
|
||||||
|
# Test with reduction 'mean'
|
||||||
|
losses_mean = nn.losses.triplet_loss(
|
||||||
|
anchors, positives, negatives, reduction="mean"
|
||||||
|
)
|
||||||
|
expected_mean = mx.mean(expected_none)
|
||||||
|
self.assertTrue(mx.allclose(losses_mean, expected_mean))
|
||||||
|
|
||||||
|
# Test with reduction 'sum'
|
||||||
|
losses_sum = nn.losses.triplet_loss(
|
||||||
|
anchors, positives, negatives, reduction="sum"
|
||||||
|
)
|
||||||
|
expected_sum = mx.sum(expected_none)
|
||||||
|
self.assertTrue(mx.allclose(losses_sum, expected_sum))
|
||||||
|
|
||||||
def test_gelu(self):
|
def test_gelu(self):
|
||||||
inputs = [1.15286231, -0.81037411, 0.35816911, 0.77484438, 0.66276414]
|
inputs = [1.15286231, -0.81037411, 0.35816911, 0.77484438, 0.66276414]
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user