Test over mr-rate/fusion-modes/vision-encoding coverage

contrastive-pretraining/mr_rate/mr_rate/mr_rate.py

@@ -2,19 +2,13 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 import torch.distributed as dist
-import torch.distributed.nn as dist_nn
 from torch.utils.checkpoint import checkpoint
-from torchvision import transforms as T, utils
-import torchvision
 
-from einops import rearrange, repeat, reduce, pack, unpack
+from einops import rearrange
 from einops.layers.torch import Rearrange
 from pathlib import Path
 import copy
-import math
-import random
 import numpy as np
-from functools import partial
 import torch.distributed.nn.functional as dist_nn_fun
 
 from transformers import BertTokenizer, BertModel

contrastive-pretraining/tests/test_fusion_modes.py

@@ -138,6 +138,13 @@ def test_cross_attn_without_text_falls_back(self, image, vol_mask):
         assert tokens.shape[1] == 18432
         assert not torch.isnan(tokens[:, :16]).any()  # real tokens
 
+        precomputed_text_latents = torch.randn(B, DIM_LATENT)
+        tokens_latents, _ = model._encode_visual_tokens(
+            image, vol_mask, vis_proj, text_latents=precomputed_text_latents,
+            num_sentences_per_image=2
+        )
+        assert tokens_latents.shape[1] == 18432
+        assert not torch.isnan(tokens[:, :16]).any() # real tokens
 
 class TestFusionConsistency:
     """All fusion modes should produce same output shape and finite values."""

contrastive-pretraining/tests/test_mr_rate_model.py

@@ -3,9 +3,10 @@
 import torch
 import torch.nn.functional as F
 import numpy as np
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 from mr_rate import MRRATE, l2norm, cast_tuple, exists
-
+from mr_rate import all_gather_batch
+from mr_rate.mr_rate import RearrangeImage
 
 # ---------------------------------------------------------------------------
 # Helper function tests
@@ -44,6 +45,40 @@ def test_cast_tuple_already_tuple(self):
     def test_cast_tuple_list(self):
         assert cast_tuple([1, 2], 2) == [1, 2]
 
+    def test_all_gather_batch_uninitialized(self):
+        x = torch.tensor([1, 2, 3])
+        assert torch.equal(all_gather_batch(x), x)
+
+    def test_initialized_distributed_concatenates_gathered_tensors(self):
+        x1 = torch.tensor([[1, 2]])
+        x2 = torch.tensor([[3, 4]])
+        with patch("torch.distributed.is_initialized", return_value=True), \
+            patch("torch.distributed.nn.functional.all_gather", return_value=[x1, x2]):
+            result = all_gather_batch(x1)
+            expected = torch.tensor([[1, 2], [3, 4]])
+
+            assert torch.equal(result, expected)
+            assert result.shape == (2, 2)
+
+    @pytest.fixture
+    def rearrange_layer(self): return RearrangeImage()
+
+    def test_rearrange_image(self, rearrange_layer):
+        x = torch.randn(2, 64, 128)
+        output = rearrange_layer(x)
+
+        assert output.ndim == 3
+        assert output.shape == (2, 64, 128)
+        assert torch.equal(output, x)
+
+    def test_rearrange_reshapes_spatial_dimensions(self):
+        b, h, w, z, c = 2, 16, 16, 4, 8
+        x_flat = torch.randn(b, h * w * z, c)
+
+        from einops import rearrange
+        expected_shape = (b, c, 16, 16, z)
+        out = rearrange(x_flat, 'b (h w z) c -> b c h w z', h=16, w=16)
+        assert out.shape == expected_shape
 
 # ---------------------------------------------------------------------------
 # MRRATE model initialization tests
@@ -123,6 +158,50 @@ def test_logit_temperature_initialized(self, mock_image_encoder, mock_text_encod
         expected = np.log(1 / 0.07)
         assert torch.allclose(model.logit_temperature, torch.tensor([expected]), atol=1e-4)
 
+    def test_init_default_text_encoder(self, mock_image_encoder, dim_text, dim_latent):
+        with patch("mr_rate.mr_rate.BertModel.from_pretrained") as mock_from_pretrained:
+            mock_bert = MagicMock()
+            mock_from_pretrained.return_value = mock_bert
+            model = MRRATE(
+                image_encoder=mock_image_encoder,
+                text_encoder=None,
+                dim_text=dim_text,
+                dim_latent=dim_latent,
+            )
+            mock_from_pretrained.assert_called_once_with("microsoft/BiomedVLP-CXR-BERT-specialized")
+            assert model.text_transformer is mock_bert
+
+    def test_init_downsample_image_embeds(self, mock_image_encoder, mock_text_encoder,
+                                        dim_text, dim_image, dim_latent):
+        from torch import nn
+        model_downsampled = MRRATE(
+            image_encoder=mock_image_encoder,
+            text_encoder=mock_text_encoder,
+            dim_text=dim_text,
+            dim_image=dim_image,
+            dim_latent=dim_latent,
+            downsample_image_embeds=True,
+        )
+        assert isinstance(model_downsampled.to_visual_latent, nn.Sequential)
+        assert isinstance(model_downsampled.to_visual_latent[0], RearrangeImage)
+        assert isinstance(model_downsampled.to_visual_latent[1], nn.Conv3d)
+        assert isinstance(model_downsampled.to_visual_latent[2], nn.Conv3d)
+
+    def test_run_checkpoint_active_branch(self, mock_image_encoder, mock_text_encoder,
+                                            dim_text, dim_image, dim_latent):
+        model = MRRATE(
+            image_encoder=mock_image_encoder,
+            text_encoder=mock_text_encoder,
+            dim_text=dim_text,
+            dim_image=dim_image,
+            dim_latent=dim_latent,
+            use_gradient_checkpointing=True
+        )
+        model.train()
+        dummy_fn = lambda x: x * 2.0
+        x = torch.randn(2, 4, requires_grad=True)
+        out = model.run_checkpoint(dummy_fn, x)
+        assert torch.allclose(out, x * 2.0)
 
 # ---------------------------------------------------------------------------
 # Forward pass tests (inference mode, no loss)
@@ -280,6 +359,19 @@ def test_loss_backward(self, model, dummy_image, dummy_text_input, real_volume_m
         grads = [p.grad for p in model.parameters() if p.grad is not None]
         assert len(grads) > 0
 
+        model.use_gradient_checkpointing = True
+        model.zero_grad()
+        loss_ckpt = model(
+            text_input=dummy_text_input,
+            image=dummy_image,
+            device='cpu',
+            real_volume_mask=real_volume_mask,
+            num_sentences_per_image=2,
+            return_loss=True,
+        )
+        loss_ckpt.backward()
+        assert torch.isfinite(loss_ckpt)
+
     def test_loss_with_sentence_mask(self, model, dummy_image, dummy_text_input, real_volume_mask):
         model.train()
         # 2 images * 2 sentences = 4 sentences, mask out last one
@@ -422,6 +514,49 @@ def test_padding_is_zero(self, model, dummy_image, real_volume_mask, dim_latent)
         padded_tokens = tokens[:, ~mask[0]]
         assert (padded_tokens == 0).all()
 
+class TestVisualInstances:
+    @pytest.fixture
+    def model(self, mock_image_encoder, mock_text_encoder, dim_text, dim_image, dim_latent):
+        return MRRATE(
+            image_encoder=mock_image_encoder,
+            text_encoder=mock_text_encoder,
+            dim_text=dim_text,
+            dim_image=dim_image,
+            dim_latent=dim_latent,
+            fusion_mode="late",
+        )
+
+    def test_instances_mask_and_shape(self, model, dummy_image, real_volume_mask):
+        model.eval()
+        vis_proj = model.to_visual_latent
+
+        tokens, mask = model._encode_visual_instances(
+            dummy_image, real_volume_mask, vis_proj
+        )
+
+        b, r = dummy_image.shape[:2]
+        tokens_per_series = 16
+
+        assert tokens.shape == (b, r * tokens_per_series, model.dim_latent), (
+            "Expected tokens shape to be == (batch_size, num_series * tokens_per_series, dim_latent)"
+        )
+        assert mask.shape == (b, r * tokens_per_series), (
+            "Expected mask shape to be == (batch_size, num_series * tokens_per_series)"
+        )
+
+        expected_mask = real_volume_mask.unsqueeze(-1).expand(b, r, tokens_per_series).reshape(b, -1)
+        assert torch.equal(mask, expected_mask), (
+            "Mask does not match expected expansion of real_volume_mask"
+        )
+        fwd_tokens, fwd_mask = model(
+            text_input=None,
+            image=dummy_image,
+            device='cpu',
+            real_volume_mask=real_volume_mask,
+            return_visual_tokens=True,
+        )
+        assert torch.equal(fwd_tokens, tokens)
+        assert torch.equal(fwd_mask, mask)
 
 # ---------------------------------------------------------------------------
 # State dict / load tests
@@ -459,7 +594,10 @@ def test_load_strips_module_prefix(self, mock_image_encoder, mock_text_encoder,
             dim_latent=dim_latent,
         )
         # Save with 'module.' prefix (simulating DDP checkpoint)
-        sd = {"module." + k: v for k, v in model.state_dict().items()}
+        sd = {}
+        for i, (k, v) in enumerate(model.state_dict().items()):
+            key = f"module.{k}" if i % 2 == 0 else k
+            sd[key] = v
         path = tmp_path / "ckpt.pt"
         torch.save(sd, path)
 
@@ -472,3 +610,5 @@ def test_load_strips_module_prefix(self, mock_image_encoder, mock_text_encoder,
         )
         model2.load(str(path))
         # Should load without error
+        for (k1, v1), (k2, v2) in zip(model.state_dict().items(), model2.state_dict().items()):
+            assert torch.equal(v1, v2)

contrastive-pretraining/tests/test_vision_encoder.py

@@ -144,7 +144,7 @@ def test_forward_cnn_mri_full_pipeline(self):
 class TestOptimizerImport:
     """Test that the optimizer utility from vision_encoder works."""
 
-    def test_get_optimizer_creates_adam(self):
+    def test_get_optimizer_creates_adamw(self):
         params = nn.Linear(10, 10).parameters()
         opt = get_optimizer(set(params), lr=1e-3, wd=0.01)
         assert opt is not None