Lifted structure loss

tensorflow_similarity/losses/__init__.py

@@ -1,4 +1,4 @@
 # Copyright 2021 The TensorFlow Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -25,5 +25,6 @@
 from .simsiam import SimSiamLoss  # noqa
 from .softnn_loss import SoftNearestNeighborLoss  # noqa
 from .triplet_loss import TripletLoss  # noqa
+from .lifted_structure_loss import LiftedStructLoss  # noqa
 from .vicreg import VicReg  # noqa
 from .xbm_loss import XBM  # noqa

tensorflow_similarity/losses/lifted_structure_loss.py

@@ -0,0 +1,121 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Lifted Structured Loss
+    Deep Metric Learning via Lifted Structured Feature Embedding.
+    https://arxiv.org/abs/1511.06452
+"""
+from __future__ import annotations
+import tensorflow as tf
+from tensorflow_similarity.algebra import build_masks
+from tensorflow_similarity.distances import Distance, distance_canonicalizer
+from tensorflow_similarity.types import FloatTensor, IntTensor
+from tensorflow_similarity import losses as tfsim_losses
+from .metric_loss import MetricLoss
+from .utils import positive_distances
+
+
+def lifted_struct_loss(
+    labels: IntTensor,
+    embeddings: FloatTensor,
+    key_labels: IntTensor,
+    key_embeddings: FloatTensor,
+    distance: Distance,
+    positive_mining_strategy: str = "hard",
+    margin: float = 1.0,
+) -> FloatTensor:
+    """Lifted Struct loss computations"""
+
+    # Compute pairwise distances
+    pairwise_distances = distance(embeddings, key_embeddings)
+
+    # Build masks for positive and negative pairs
+    positive_mask, negative_mask = build_masks(
+        query_labels=labels, key_labels=key_labels, batch_size=tf.shape(embeddings)[0]
+    )
+
+    # Get positive distances and indices
+    positive_dists, positive_indices = positive_distances(positive_mining_strategy, pairwise_distances, positive_mask)
+
+    # Reorder pairwise distances and negative mask based on positive indices
+    reordered_pairwise_distances = tf.gather(pairwise_distances, positive_indices, axis=1)
+    reordered_negative_mask = tf.gather(negative_mask, positive_indices, axis=1)
+
+    # Concatenate pairwise distances and negative masks along axis=1
+    concatenated_distances = tf.concat([pairwise_distances, reordered_pairwise_distances], axis=1)
+    concatenated_negative_mask = tf.concat([negative_mask, reordered_negative_mask], axis=1)
+    concatenated_negative_mask = tf.cast(concatenated_negative_mask, tf.float32)
+    # Compute (margin - neg_dist) logsum_exp values for each row (equation 4 in the paper)
+    neg_logsumexp = tfsim_losses.utils.logsumexp(margin - concatenated_distances, concatenated_negative_mask)
+
+    # Calculate the loss
+    j_values = neg_logsumexp + positive_dists
+
+    loss: FloatTensor = j_values / 2.0
+
+    return loss
+
+
+@tf.keras.utils.register_keras_serializable(package="Similarity")
+class LiftedStructLoss(MetricLoss):
+    """Computes the lifted structured loss in an online fashion.
+    This loss encourages the positive distances between a pair of embeddings
+    with the same labels to be smaller than the negative distances between pair
+    of embeddings of different labels.
+    See: https://arxiv.org/abs/1511.06452 for the original paper.
+    `y_true` must be a 1-D integer `Tensor` of shape (batch_size,).
+    It's values represent the classes associated with the examples as
+    **integer values**.
+    `y_pred` must be 2-D float `Tensor` of L2 normalized embedding vectors.
+    You can use the layer `tensorflow_similarity.layers.L2Embedding()` as the
+    last layer of your model to ensure your model output is properly normalized.
+    """
+
+    def __init__(
+        self,
+        distance: Distance | str = "cosine",
+        positive_mining_strategy: str = "hard",
+        margin: float = 1.0,
+        name: str = "LiftedStructLoss",
+        **kwargs,
+    ):
+        """Initializes the LiftedStructLoss.
+        Args:
+            distance: Which distance function to use to compute the pairwise
+                distances between embeddings.
+            positive_mining_strategy: What mining strategy to use to select
+                embedding from the same class. Defaults to 'hard'.
+                Available: {'easy', 'hard'}
+            margin: Use an explicit value for the margin term.
+            name: Loss name. Defaults to "LiftedStructLoss".
+        Raises:
+            ValueError: Invalid positive mining strategy.
+        """
+
+        # distance canonicalization
+        distance = distance_canonicalizer(distance)
+        self.distance = distance
+
+        # sanity checks
+        if positive_mining_strategy not in ["easy", "hard"]:
+            raise ValueError("Invalid positive mining strategy")
+
+        super().__init__(
+            lifted_struct_loss,
+            name=name,
+            distance=distance,
+            positive_mining_strategy=positive_mining_strategy,
+            margin=margin,
+            **kwargs,
+        )

tests/losses/test_lifted_structure_loss.py

@@ -0,0 +1,65 @@
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow.python.framework import combinations
+from tensorflow.keras.losses import Reduction
+from tensorflow_similarity import losses
+from . import utils
+
+
+@combinations.generate(combinations.combine(mode=["graph", "eager"]))
+class TestLiftedStructLoss(tf.test.TestCase, parameterized.TestCase):
+    def test_config(self):
+        lifted_obj = losses.LiftedStructLoss(
+            reduction=Reduction.SUM,
+            name="lifted_loss",
+        )
+        self.assertEqual(lifted_obj.distance.name, "cosine")
+        self.assertEqual(lifted_obj.name, "lifted_loss")
+        self.assertEqual(lifted_obj.reduction, Reduction.SUM)
+
+    @parameterized.named_parameters(
+        {"testcase_name": "_fixed_margin", "margin": 1.1, "expected_loss": 157.68167},
+    )
+    def test_all_correct_unweighted(self, margin, expected_loss):
+        """Tests the LiftedStructLoss with different parameters."""
+        y_true, y_preds = utils.generate_perfect_test_batch()
+
+        lifted_obj = losses.LiftedStructLoss(reduction=Reduction.SUM, margin=margin)
+        loss = lifted_obj(y_true, y_preds)
+        self.assertAlmostEqual(self.evaluate(loss), expected_loss, 3)
+
+    @parameterized.named_parameters(
+        {"testcase_name": "_fixed_margin", "margin": 1.0, "expected_loss": 187.37393},
+    )
+    def test_all_mismatch_unweighted(self, margin, expected_loss):
+        """Tests the LiftedStructLoss with different parameters."""
+        y_true, y_preds = utils.generate_bad_test_batch()
+
+        lifted_obj = losses.LiftedStructLoss(reduction=Reduction.SUM, margin=margin)
+        loss = lifted_obj(y_true, y_preds)
+        self.assertAlmostEqual(self.evaluate(loss), expected_loss, 3)
+
+    @parameterized.named_parameters(
+        {"testcase_name": "_fixed_margin", "margin": 1.0, "expected_loss": 2.927718},
+    )
+    def test_no_reduction(self, margin, expected_loss):
+        """Tests the LiftedStructLoss with different parameters."""
+        y_true, y_preds = utils.generate_bad_test_batch()
+
+        lifted_obj = losses.LiftedStructLoss(reduction=Reduction.NONE, margin=margin)
+        loss = lifted_obj(y_true, y_preds)
+        loss = self.evaluate(loss)
+        expected_loss = self.evaluate(tf.fill(y_true.shape, expected_loss))
+        self.assertArrayNear(loss, expected_loss, 0.001)
+
+    @parameterized.named_parameters(
+        {"testcase_name": "_fixed_margin", "margin": 1.0, "expected_loss": 2.414156913757324},
+    )
+    def test_sum_reduction(self, margin, expected_loss):
+        """Tests the LiftedStructLoss with different parameters."""
+        y_true, y_preds = utils.generate_perfect_test_batch()
+
+        lifted_obj = losses.LiftedStructLoss(reduction=Reduction.SUM, margin=margin)
+        loss = lifted_obj(y_true, y_preds)
+        expected_loss = y_true.shape[0] * expected_loss
+        self.assertAlmostEqual(self.evaluate(loss), expected_loss, 3)