diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py index d8b2c33e62..91225fc61a 100644 --- a/monai/losses/contrastive.py +++ b/monai/losses/contrastive.py @@ -28,20 +28,38 @@ class ContrastiveLoss(_Loss): Adapted from: https://github.com/Sara-Ahmed/SiT/blob/1aacd6adcd39b71efc903d16b4e9095b97dda76f/losses.py#L5 + The loss is evaluated using positive-shifted logits in log space. By default, float16 + and bfloat16 inputs use float32 arithmetic to preserve precision and accumulation range. + + Very small temperatures can amplify roundoff in cosine-similarity gradients. + For higher precision, pass both embeddings as float64 before evaluating the loss, + for example ``loss(input.double(), target.double())``, or set ``compute_dtype=torch.float64``. + Casting the computed loss to float64 does not improve the precision of its backward computation. + """ - def __init__(self, temperature: float = 0.5, batch_size: int = -1) -> None: + def __init__( + self, temperature: float = 0.5, batch_size: int = -1, *, compute_dtype: torch.dtype | None = None + ) -> None: """ Args: temperature: Can be scaled between 0 and 1 for learning from negative samples, ideally set to 0.5. + batch_size: Legacy parameter. The batch size is estimated dynamically in ``forward``. + compute_dtype: Floating-point dtype for similarity and loss computations, either ``torch.float32`` + or ``torch.float64``. Defaults to None, preserving float64 inputs and promoting float16/bfloat16 + inputs to float32. Selecting float64 can reduce backward roundoff at very small temperatures + on devices supporting float64, with additional computation and memory costs. Gradients are + still stored in the original input dtype; existing input quantization cannot be recovered. Raises: - ValueError: When an input of dimension length > 2 is passed - ValueError: When input and target are of different shapes + ValueError: If ``compute_dtype`` is neither None nor a supported computation dtype. """ super().__init__() + if compute_dtype not in (None, torch.float32, torch.float64): + raise ValueError("compute_dtype must be None, torch.float32, or torch.float64.") self.temperature = temperature + self.compute_dtype = compute_dtype if batch_size != -1: warn("batch_size is no longer required to be set. It will be estimated dynamically in the forward call") @@ -51,6 +69,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: Args: input: the shape should be B[F]. target: the shape should be B[F]. + + Returns: + A scalar loss tensor in ``compute_dtype`` when specified. Otherwise, float16/bfloat16 + inputs use float32, and float32/float64 inputs retain their dtype. + + Raises: + ValueError: If either input has more than two dimensions or their shapes differ. """ if len(target.shape) > 2 or len(input.shape) > 2: raise ValueError( @@ -61,21 +86,25 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if target.shape != input.shape: raise ValueError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})") - temperature_tensor = torch.as_tensor(self.temperature).to(input.device) batch_size = input.shape[0] - negatives_mask = ~torch.eye(batch_size * 2, batch_size * 2, dtype=torch.bool) - negatives_mask = torch.clone(negatives_mask.type(torch.float)).to(input.device) + self_mask = torch.eye(batch_size * 2, dtype=torch.bool, device=input.device) repr = torch.cat([input, target], dim=0) + compute_dtype = self.compute_dtype + if compute_dtype is not None: + repr = repr.to(dtype=compute_dtype) + elif repr.dtype in (torch.float16, torch.bfloat16): + repr = repr.float() + temperature_tensor = torch.as_tensor(self.temperature, dtype=repr.dtype, device=input.device) sim_matrix = F.cosine_similarity(repr.unsqueeze(1), repr.unsqueeze(0), dim=2) sim_ij = torch.diag(sim_matrix, batch_size) sim_ji = torch.diag(sim_matrix, -batch_size) positives = torch.cat([sim_ij, sim_ji], dim=0) - nominator = torch.exp(positives / temperature_tensor) - denominator = negatives_mask * torch.exp(sim_matrix / temperature_tensor) + # Shift before scaling so the positive logit is zero, avoiding cancellation after logsumexp. + logits = ((sim_matrix - positives.unsqueeze(1)) / temperature_tensor).masked_fill(self_mask, -torch.inf) - loss_partial = -torch.log(nominator / torch.sum(denominator, dim=1)) + loss_partial = torch.logsumexp(logits, dim=1) return torch.sum(loss_partial) / (2 * batch_size) diff --git a/tests/losses/test_contrastive_loss.py b/tests/losses/test_contrastive_loss.py index 21a9e76417..325d9cd311 100644 --- a/tests/losses/test_contrastive_loss.py +++ b/tests/losses/test_contrastive_loss.py @@ -11,6 +11,7 @@ from __future__ import annotations +import math import unittest import numpy as np @@ -56,6 +57,125 @@ class TestContrastiveLoss(unittest.TestCase): + @parameterized.expand( + [(temperature, dtype) for temperature in (0.5, 0.01, 0.001) for dtype in (torch.float16, torch.bfloat16)] + ) + def test_reduced_precision_loss_and_gradients(self, temperature, dtype): + # All nonself similarities are 0.5, but the embedding gradients are nonzero. + vectors = torch.cat([torch.ones(4, 1), torch.eye(4)], dim=1).to(dtype) + inputs = vectors[:2].clone().requires_grad_() + targets = vectors[2:].clone().requires_grad_() + ref_inputs = inputs.detach().double().requires_grad_() + ref_targets = targets.detach().double().requires_grad_() + representations = torch.nn.functional.normalize(torch.cat([ref_inputs, ref_targets]), dim=1) + logits = (representations @ representations.T / temperature).masked_fill( + torch.eye(4, dtype=torch.bool), -torch.inf + ) + expected = torch.nn.functional.cross_entropy(logits, torch.tensor([2, 3, 0, 1])) + result = ContrastiveLoss(temperature=temperature)(inputs, targets) + self.assertEqual(result.dtype, torch.float32) + torch.testing.assert_close(result.double(), expected, atol=1e-5, rtol=1e-5) + result.backward() + expected.backward() + torch.testing.assert_close(inputs.grad, ref_inputs.grad.to(dtype), atol=1e-4, rtol=1e-2) + torch.testing.assert_close(targets.grad, ref_targets.grad.to(dtype), atol=1e-4, rtol=1e-2) + + @parameterized.expand( + [(1e-7, torch.float32), (1e-8, torch.float32), (1e-12, torch.float64), (1e-50, torch.float64)] + ) + def test_small_temperature_cancellation(self, temperature, dtype): + inputs = torch.ones((2, 4), dtype=dtype, requires_grad=True) + targets = torch.ones_like(inputs, requires_grad=True) + result = ContrastiveLoss(temperature=temperature)(inputs, targets) + torch.testing.assert_close(result, result.new_tensor(math.log(3)), atol=1e-7, rtol=1e-7) + result.backward() + self.assertTrue(torch.isfinite(inputs.grad).all()) + self.assertTrue(torch.isfinite(targets.grad).all()) + + def test_float16_batch_reduction(self): + # The mean is finite although a float16 sum of the per-anchor losses overflows. + batch_size, temperature = 1280, 0.1 + inputs = torch.tensor([[1.0, 0.0]], dtype=torch.float16).repeat(batch_size, 1) + result = ContrastiveLoss(temperature=temperature)(inputs, -inputs) + expected = 2 / temperature + math.log(batch_size - 1 + batch_size * math.exp(-2 / temperature)) + self.assertEqual(result.dtype, torch.float32) + torch.testing.assert_close(result, result.new_tensor(expected)) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_cpu_autocast(self, dtype): + vectors = torch.cat([torch.ones(4, 1), torch.eye(4)], dim=1).to(dtype) + inputs = vectors[:2].clone().requires_grad_() + targets = vectors[2:].clone().requires_grad_() + with torch.autocast(device_type="cpu", dtype=dtype): + result = ContrastiveLoss(temperature=0.001)(inputs, targets) + self.assertEqual(result.dtype, torch.float32) + torch.testing.assert_close(result, result.new_tensor(math.log(3))) + result.backward() + self.assertTrue(torch.isfinite(inputs.grad).all()) + self.assertTrue(torch.isfinite(targets.grad).all()) + + def test_gradcheck(self): + inputs = torch.tensor([[1.0, 0.2, -0.3], [-0.1, 0.7, 0.5]], dtype=torch.float64, requires_grad=True) + targets = torch.tensor([[0.8, -0.4, 0.1], [0.3, 0.6, -0.2]], dtype=torch.float64, requires_grad=True) + self.assertTrue(torch.autograd.gradcheck(ContrastiveLoss(temperature=0.1), (inputs, targets))) + + @parameterized.expand( + [(temperature, dtype) for temperature in (0.01, 0.001) for dtype in (torch.float32, torch.float64)] + ) + def test_identical_low_temperature(self, temperature, dtype): + inputs = torch.ones((2, 4), dtype=dtype, requires_grad=True) + targets = torch.ones((2, 4), dtype=dtype, requires_grad=True) + result = ContrastiveLoss(temperature=temperature)(inputs, targets) + # Every anchor has three equally similar candidates, including its positive. + torch.testing.assert_close(result, result.new_tensor(math.log(3)), atol=1e-4, rtol=1e-4) + result.backward() + self.assertTrue(torch.isfinite(inputs.grad).all()) + self.assertTrue(torch.isfinite(targets.grad).all()) + + @parameterized.expand( + [(temperature, dtype) for temperature in (0.01, 0.001) for dtype in (torch.float32, torch.float64)] + ) + def test_antipodal_low_temperature(self, temperature, dtype): + inputs = torch.eye(2, dtype=dtype, requires_grad=True) + targets = (-torch.eye(2, dtype=dtype)).requires_grad_() + result = ContrastiveLoss(temperature=temperature)(inputs, targets) + # Positive similarity is -1, and both other candidates have similarity 0. + expected = 1 / temperature + math.log(2 + math.exp(-1 / temperature)) + torch.testing.assert_close(result, result.new_tensor(expected)) + result.backward() + self.assertTrue(torch.isfinite(inputs.grad).all()) + self.assertTrue(torch.isfinite(targets.grad).all()) + + @parameterized.expand( + [(temperature, dtype) for temperature in (0.5, 0.05, 0.01) for dtype in (torch.float32, torch.float64)] + ) + def test_reference_loss_and_gradients(self, temperature, dtype): + inputs = torch.tensor([[1.0, 0.2, -0.3], [-0.1, 0.7, 0.5]], dtype=dtype, requires_grad=True) + targets = torch.tensor([[0.8, -0.4, 0.1], [0.3, 0.6, -0.2]], dtype=dtype, requires_grad=True) + reference_inputs = inputs.detach().double().requires_grad_() + reference_targets = targets.detach().double().requires_grad_() + representations = torch.nn.functional.normalize(torch.cat([reference_inputs, reference_targets]), dim=1) + logits = representations @ representations.T / temperature + logits = logits.masked_fill(torch.eye(4, dtype=torch.bool), -torch.inf) + # The positive for each anchor is the corresponding sample in the other view. + expected = torch.nn.functional.cross_entropy(logits, torch.tensor([2, 3, 0, 1])) + result = ContrastiveLoss(temperature=temperature)(inputs, targets) + torch.testing.assert_close(result.double(), expected, atol=1e-4, rtol=1e-4) + result.backward() + expected.backward() + torch.testing.assert_close(inputs.grad.double(), reference_inputs.grad, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(targets.grad.double(), reference_targets.grad, atol=1e-4, rtol=1e-4) + + @parameterized.expand([(temperature, sign) for temperature in (0.01, 0.001) for sign in (1, -1)]) + def test_single_pair_low_temperature(self, temperature, sign): + inputs = torch.tensor([[1.0, 0.0]], requires_grad=True) + targets = (sign * inputs.detach()).requires_grad_() + result = ContrastiveLoss(temperature=temperature)(inputs, targets) + torch.testing.assert_close(result, torch.zeros_like(result)) + result.backward() + torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs)) + torch.testing.assert_close(targets.grad, torch.zeros_like(targets)) + @parameterized.expand(TEST_CASES) def test_result(self, input_param, input_data, expected_val): contrastiveloss = ContrastiveLoss(**input_param) @@ -82,5 +202,109 @@ def check_warning_rasied(self): ContrastiveLoss(temperature=0.5, batch_size=1) +class TestContrastiveComputeDtype(unittest.TestCase): + + @parameterized.expand([(features, seed) for features in (3, 5, 17) for seed in range(8)]) + def test_identical_extreme_temperature_gradients(self, features, seed): + generator = torch.Generator().manual_seed(seed) + vector = torch.randn((1, features), generator=generator) + inputs = vector.repeat(2, 1).requires_grad_() + targets = vector.repeat(2, 1).requires_grad_() + result = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)(inputs, targets) + self.assertEqual(result.dtype, torch.float64) + torch.testing.assert_close(result, result.new_tensor(math.log(3))) + result.backward() + # Each cosine is stationary at identical nonzero vectors; its mathematical gradient is zero. + torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs), atol=1e-7, rtol=0) + torch.testing.assert_close(targets.grad, torch.zeros_like(targets), atol=1e-7, rtol=0) + + @parameterized.expand( + [ + (dtype, temperature) + for dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64) + for temperature in (0.5, 0.001) + ] + + [(torch.float32, 1e-8), (torch.float64, 1e-8)] + ) + def test_equivalent_to_casting_inputs(self, dtype, temperature): + vectors = torch.cat([torch.ones(4, 1), torch.eye(4)], dim=1).to(dtype) + inputs = vectors[:2].clone().requires_grad_() + targets = vectors[2:].clone().requires_grad_() + original_inputs, original_targets = inputs.detach().clone(), targets.detach().clone() + ref_inputs = inputs.detach().clone().requires_grad_() + ref_targets = targets.detach().clone().requires_grad_() + result = ContrastiveLoss(temperature=temperature, compute_dtype=torch.float64)(inputs, targets) + expected = ContrastiveLoss(temperature=temperature)(ref_inputs.double(), ref_targets.double()) + torch.testing.assert_close(result, expected, atol=0, rtol=0) + result.backward() + expected.backward() + torch.testing.assert_close(inputs.grad, ref_inputs.grad, atol=0, rtol=0) + torch.testing.assert_close(targets.grad, ref_targets.grad, atol=0, rtol=0) + torch.testing.assert_close(inputs, original_inputs, atol=0, rtol=0) + torch.testing.assert_close(targets, original_targets, atol=0, rtol=0) + self.assertTrue(torch.isfinite(inputs.grad).all()) + self.assertTrue(torch.isfinite(targets.grad).all()) + + @parameterized.expand( + [ + (input_dtype, compute_dtype) + for input_dtype in (torch.float32, torch.float64) + for compute_dtype in (None, torch.float32, torch.float64) + ] + ) + def test_output_dtype(self, input_dtype, compute_dtype): + inputs = torch.ones((2, 3), dtype=input_dtype) + result = ContrastiveLoss(compute_dtype=compute_dtype)(inputs, inputs) + self.assertEqual(result.dtype, input_dtype if compute_dtype is None else compute_dtype) + + @parameterized.expand([(None,), (torch.float32,), (torch.float64,)]) + def test_mixed_input_dtype_backward(self, compute_dtype): + inputs = torch.tensor([[1.0, 0.2, -0.3], [-0.1, 0.7, 0.5]], requires_grad=True) + targets = torch.tensor([[0.8, -0.4, 0.1], [0.3, 0.6, -0.2]], dtype=torch.float64, requires_grad=True) + ref_inputs = inputs.detach().clone().requires_grad_() + ref_targets = targets.detach().clone().requires_grad_() + result = ContrastiveLoss(temperature=0.001, compute_dtype=compute_dtype)(inputs, targets) + expected_dtype = torch.float64 if compute_dtype is None else compute_dtype + expected = ContrastiveLoss(temperature=0.001)(ref_inputs.to(expected_dtype), ref_targets.to(expected_dtype)) + torch.testing.assert_close(result, expected, atol=0, rtol=0) + result.backward() + expected.backward() + torch.testing.assert_close(inputs.grad, ref_inputs.grad, atol=0, rtol=0) + torch.testing.assert_close(targets.grad, ref_targets.grad, atol=0, rtol=0) + self.assertEqual(inputs.grad.dtype, torch.float32) + self.assertEqual(targets.grad.dtype, torch.float64) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,), (torch.int64,), (torch.bool,), ("float64",)]) + def test_invalid_compute_dtype(self, compute_dtype): + with self.assertRaisesRegex(ValueError, "compute_dtype"): + ContrastiveLoss(compute_dtype=compute_dtype) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_cpu_autocast_respects_compute_dtype(self, autocast_dtype): + inputs = torch.ones((2, 3), requires_grad=True) + targets = torch.ones_like(inputs, requires_grad=True) + with torch.autocast(device_type="cpu", dtype=autocast_dtype): + result = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)(inputs, targets) + self.assertEqual(result.dtype, torch.float64) + result.backward() + torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs), atol=1e-7, rtol=0) + torch.testing.assert_close(targets.grad, torch.zeros_like(targets), atol=1e-7, rtol=0) + + def test_singleton_high_precision(self): + inputs = torch.tensor([[0.3, -0.7, 0.2]], requires_grad=True) + targets = torch.tensor([[-0.4, 0.1, 0.8]], requires_grad=True) + result = ContrastiveLoss(temperature=1e-8, compute_dtype=torch.float64)(inputs, targets) + torch.testing.assert_close(result, torch.zeros_like(result)) + result.backward() + torch.testing.assert_close(inputs.grad, torch.zeros_like(inputs)) + torch.testing.assert_close(targets.grad, torch.zeros_like(targets)) + + def test_legacy_arguments_preserved(self): + with self.assertWarnsRegex(Warning, "batch_size"): + loss = ContrastiveLoss(0.5, 2, compute_dtype=torch.float64) + inputs = torch.ones((2, 3)) + self.assertEqual(loss(inputs, inputs).dtype, torch.float64) + + if __name__ == "__main__": unittest.main()