[Triplet Margin Loss] Issue 1118 #1120
Conversation
|
@vroulet May I know if there's anything that needs to be changed? |
| anchor: The anchor embeddings. Shape: [batch_size, feature_dim]. | ||
| positive: The positive embeddings. Shape: [batch_size, feature_dim]. | ||
| negative: The negative embeddings. Shape: [batch_size, feature_dim]. | ||
| margin: The margin value. Default: 1.0. |
There was a problem hiding this comment.
No need to put the default values since they are given in the signature.
| by V. Balntas et al. Default: False. | ||
| reduction: Specifies the reduction to apply to the output: | ||
| 'none' | 'mean' | 'sum'. Default: 'mean'. | ||
|
|
| margin: The margin value. Default: 1.0. | ||
| p: The norm degree for pairwise distance. Default: 2. | ||
| eps: Small epsilon value to avoid numerical issues. Default: 1e-6. | ||
| swap: Use the distance swap optimization from "Learning shallow |
There was a problem hiding this comment.
Use rst formatting for references (see e.g. the docstring of Adam)
| swap: bool = False, | ||
| reduction: str = 'mean', | ||
| ) -> chex.Array: | ||
| """Triplet margin loss function. |
| @@ -53,5 +53,41 @@ def test_batched(self): | |||
| ) | |||
|
|
|||
|
|
|||
| class TripletMarginLossTest(chex.TestCase): | |||
|
|
|||
| def setUp(self): | |||
There was a problem hiding this comment.
Avoid using numerical values as expected returns.
They may fail depending on the backend for example.
You may consider simple test cases with a "handmade" function (see e.g. the lbfgs tests). You can check for specific inputs (like zeros or ones).
You may also add a test for some specific behaviors (like using swap here).
Also you should test this function under jit/vmap etc... (see the chex.all_variant utility in some other tests).
@vroulet Hi Vincent, Added code and tests for the Triplet Margin Loss Function. Kindly review the code and please do comment in case of any changes.