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 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 44ffc22..bbef7ff 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,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) 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