diff --git a/monai/networks/layers/vector_quantizer.py b/monai/networks/layers/vector_quantizer.py index 388f93fe2df..30571ed652d 100644 --- a/monai/networks/layers/vector_quantizer.py +++ b/monai/networks/layers/vector_quantizer.py @@ -105,6 +105,8 @@ def quantize(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, to del encoding_indices_view[1] inputs = inputs.float() + # The codebook is upcast too, so a model cast to float16/bfloat16 (not just autocast) still works. + embedding_weight = self.embedding.weight.float() # Converting to channel last format flat_input = inputs.permute(self.flatten_permutation).contiguous().view(-1, self.embedding_dim) @@ -112,8 +114,8 @@ def quantize(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, to # Calculate Euclidean distances distances = ( (flat_input**2).sum(dim=1, keepdim=True) - + (self.embedding.weight.t() ** 2).sum(dim=0, keepdim=True) - - 2 * torch.mm(flat_input, self.embedding.weight.t()) + + (embedding_weight.t() ** 2).sum(dim=0, keepdim=True) + - 2 * torch.mm(flat_input, embedding_weight.t()) ) # Mapping distances to indexes diff --git a/tests/networks/layers/test_vector_quantizer.py b/tests/networks/layers/test_vector_quantizer.py index 43533d03771..54719642f66 100644 --- a/tests/networks/layers/test_vector_quantizer.py +++ b/tests/networks/layers/test_vector_quantizer.py @@ -68,6 +68,19 @@ def test_ema(self): self.assertTrue(all(layer.embedding.weight[0] != original_weight_0)) self.assertTrue(all(layer.embedding.weight[1] == original_weight_1)) + @parameterized.expand([[torch.float16], [torch.bfloat16]]) + def test_ema_half_precision_module(self, dtype): + """A quantizer cast to half precision (not just autocast) runs and picks the same codes as float32.""" + layer = EMAQuantizer(spatial_dims=2, num_embeddings=16, embedding_dim=8).eval() + x = torch.randn(1, 8, 4, 4) + expected_indices = layer.quantize(x)[2] + + layer = layer.to(dtype) + quantized, _, encoding_indices = layer(x.to(dtype)) + + self.assertEqual(quantized.dtype, dtype) + torch.testing.assert_close(encoding_indices, expected_indices) + class TestVectorQuantizer(unittest.TestCase): @parameterized.expand(TEST_CASES)