From 99b8e2e0382b5090e3664af6e71de34227b6e1f0 Mon Sep 17 00:00:00 2001 From: zephyrus Date: Sat, 15 Aug 2026 03:32:32 +0300 Subject: [PATCH 1/5] refactor(mr_rate): removed unused imports to optimize minor at startup --- contrastive-pretraining/mr_rate/mr_rate/mr_rate.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py b/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py index 008b129..4bc5c29 100644 --- a/contrastive-pretraining/mr_rate/mr_rate/mr_rate.py +++ b/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 From 9e09d3c14f7552d9593844bdb1c6aab05a8c9c55 Mon Sep 17 00:00:00 2001 From: zephyrus Date: Fri, 21 Aug 2026 19:56:29 +0300 Subject: [PATCH 2/5] test(model): expand unit test coverage for MRRATE and utilities - Add tests for all_gather_batch in single-device and distributed modes - Add tests for RearrangeImage and downsampled visual latent projections - Add default BiomedVLP text encoder initialization test - Add TestVisualInstances for instance encoding and mask expansion - Assert state dict weight equivalence in model loading test --- .../tests/test_mr_rate_model.py | 109 +++++++++++++++++- 1 file changed, 106 insertions(+), 3 deletions(-) diff --git a/contrastive-pretraining/tests/test_mr_rate_model.py b/contrastive-pretraining/tests/test_mr_rate_model.py index 44ffc22..77e7976 100644 --- a/contrastive-pretraining/tests/test_mr_rate_model.py +++ b/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,35 @@ 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) + # --------------------------------------------------------------------------- # Forward pass tests (inference mode, no loss) @@ -422,6 +486,40 @@ 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" + ) # --------------------------------------------------------------------------- # State dict / load tests @@ -459,7 +557,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 +573,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) From 0ef69344bfe64edd3b30b5f54e3038d91a793695 Mon Sep 17 00:00:00 2001 From: Zephyruss1 Date: Mon, 24 Aug 2026 21:59:57 +0300 Subject: [PATCH 3/5] test(mr-rate/fusion-modes): improve test coverage for checkpointing and visual tokens - Add test for run_checkpoint active branch during training mode - Add gradient checkpointing backward pass test in TestMRRATELoss - Add test for forward call with return_visual_tokens=True - Add test for precomputed 2D text_latents in late attention pooling --- .../tests/test_fusion_modes.py | 7 ++++ .../tests/test_mr_rate_model.py | 37 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/contrastive-pretraining/tests/test_fusion_modes.py b/contrastive-pretraining/tests/test_fusion_modes.py index 1d2a68c..2045745 100644 --- a/contrastive-pretraining/tests/test_fusion_modes.py +++ b/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.""" diff --git a/contrastive-pretraining/tests/test_mr_rate_model.py b/contrastive-pretraining/tests/test_mr_rate_model.py index 77e7976..915e561 100644 --- a/contrastive-pretraining/tests/test_mr_rate_model.py +++ b/contrastive-pretraining/tests/test_mr_rate_model.py @@ -187,6 +187,21 @@ def test_init_downsample_image_embeds(self, mock_image_encoder, mock_text_encode 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) @@ -344,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 @@ -520,6 +548,15 @@ def test_instances_mask_and_shape(self, model, dummy_image, real_volume_mask): 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 From 8b89173ae4f9259a0042627e41310e2e95299a29 Mon Sep 17 00:00:00 2001 From: Zephyruss1 Date: Mon, 24 Aug 2026 22:04:07 +0300 Subject: [PATCH 4/5] test(vision-encoder): rename optimizer test to reflect AdamW creation - Since get_optimizer returns AdamW when weight decay > 0 (in this case wd=0.01), this accurately reflects the optimizer type being verified. --- contrastive-pretraining/tests/test_vision_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contrastive-pretraining/tests/test_vision_encoder.py b/contrastive-pretraining/tests/test_vision_encoder.py index 16f866b..b40f43c 100644 --- a/contrastive-pretraining/tests/test_vision_encoder.py +++ b/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 From 12570a7d857c43484ff4bf475c3720fb7349e342 Mon Sep 17 00:00:00 2001 From: zephyrus Date: Wed, 26 Aug 2026 01:31:12 +0300 Subject: [PATCH 5/5] test(mr-rate): minor indent fix --- .../tests/test_mr_rate_model.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/contrastive-pretraining/tests/test_mr_rate_model.py b/contrastive-pretraining/tests/test_mr_rate_model.py index 915e561..bbef7ff 100644 --- a/contrastive-pretraining/tests/test_mr_rate_model.py +++ b/contrastive-pretraining/tests/test_mr_rate_model.py @@ -189,19 +189,19 @@ def test_init_downsample_image_embeds(self, mock_image_encoder, mock_text_encode 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) + 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)