Repository: Project-MONAI/MONAI

Issue:
ContrastiveLoss with dynamic batchsize
**Is your feature request related to a problem? Please describe.**
When we instantiate an instance of type class `ContrastiveLoss` we need to pass in the parameter for `batch_size`. This assumes that the batch size remains constant all through the training(which is not always the case). 

**Describe the solution you'd like**
As we know and have fixed in the [docs](https://docs.monai.io/en/latest/losses.html#monai.losses.ContrastiveLoss) that the forward accepts inputs of shape `B[F]`, why not pick up the Batch size dynamically in the forward. Is there a reason why it is not designed in that fashion?

I would be happy to put a pull request if it is okay.



Discussion:
@a-parida12 Please go ahead and send a PR, I feel we should make the dynamic part optional so that we don't break existing code for other users and keep the default dynamic_batch_size=False

We need insights from @wyli if it's okay for the optional idea :)
Thanks, @finalelement! I thought the majority of the losses from PyTorch don't require that the batch size be known and fixed so extending that to this too. But you are right we need to be backward compatible too.  

How about I raise a "warning" when the `batch_size` parameter is set to let users know it is unnecessary ?

Also will wait for @wyli 's input and then begin the PR.
Also tagging this discussion for the completeness of my motivation https://github.com/Project-MONAI/tutorials/discussions/624#discussion-3957066
yes please go ahead, would be great to automatically infer the batch size whenever possible.

Patch:
diff --git a/monai/losses/contrastive.py b/monai/losses/contrastive.py
--- a/monai/losses/contrastive.py
+++ b/monai/losses/contrastive.py
@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from distutils.log import warn
+
 import torch
 from torch.nn import functional as F
 from torch.nn.modules.loss import _Loss
@@ -30,11 +32,10 @@ class ContrastiveLoss(_Loss):
     """
 
     @deprecated_arg(name="reduction", since="0.8", msg_suffix="`reduction` is no longer supported.")
-    def __init__(self, temperature: float = 0.5, batch_size: int = 1, reduction="sum") -> None:
+    def __init__(self, temperature: float = 0.5, batch_size: int = -1, reduction="sum") -> None:
         """
         Args:
             temperature: Can be scaled between 0 and 1 for learning from negative samples, ideally set to 0.5.
-            batch_size: The number of samples.
 
         Raises:
             ValueError: When an input of dimension length > 2 is passed
@@ -46,10 +47,11 @@ def __init__(self, temperature: float = 0.5, batch_size: int = 1, reduction="sum
 
         """
         super().__init__()
-
-        self.batch_size = batch_size
         self.temperature = temperature
 
+        if batch_size != -1:
+            warn("batch_size is no longer required to be set. It will be estimated dynamically in the forward call")
+
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         Args:
@@ -66,17 +68,18 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             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]
 
         norm_i = F.normalize(input, dim=1)
         norm_j = F.normalize(target, dim=1)
 
-        negatives_mask = ~torch.eye(self.batch_size * 2, self.batch_size * 2, dtype=torch.bool)
+        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)
 
         repr = torch.cat([norm_i, norm_j], dim=0)
         sim_matrix = F.cosine_similarity(repr.unsqueeze(1), repr.unsqueeze(0), dim=2)
-        sim_ij = torch.diag(sim_matrix, self.batch_size)
-        sim_ji = torch.diag(sim_matrix, -self.batch_size)
+        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)
@@ -84,4 +87,4 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
         loss_partial = -torch.log(nominator / torch.sum(denominator, dim=1))
 
-        return torch.sum(loss_partial) / (2 * self.batch_size)
+        return torch.sum(loss_partial) / (2 * batch_size)


Test patch:
diff --git a/tests/test_contrastive_loss.py b/tests/test_contrastive_loss.py
--- a/tests/test_contrastive_loss.py
+++ b/tests/test_contrastive_loss.py
@@ -19,12 +19,12 @@
 
 TEST_CASES = [
     [  # shape: (1, 4), (1, 4)
-        {"temperature": 0.5, "batch_size": 1},
+        {"temperature": 0.5},
         {"input": torch.tensor([[1.0, 1.0, 0.0, 0.0]]), "target": torch.tensor([[1.0, 1.0, 0.0, 0.0]])},
         0.0,
     ],
     [  # shape: (2, 4), (2, 4)
-        {"temperature": 0.5, "batch_size": 2},
+        {"temperature": 0.5},
         {
             "input": torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]),
             "target": torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]),
@@ -32,7 +32,7 @@
         1.0986,
     ],
     [  # shape: (1, 4), (1, 4)
-        {"temperature": 0.5, "batch_size": 2},
+        {"temperature": 0.5},
         {
             "input": torch.tensor([[1.0, 2.0, 3.0, 4.0], [1.0, 1.0, 0.0, 0.0]]),
             "target": torch.tensor([[0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]),
@@ -40,12 +40,12 @@
         0.8719,
     ],
     [  # shape: (1, 4), (1, 4)
-        {"temperature": 0.5, "batch_size": 1},
+        {"temperature": 0.5},
         {"input": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), "target": torch.tensor([[1.0, 1.0, 0.0, 0.0]])},
         0.0,
     ],
     [  # shape: (1, 4), (1, 4)
-        {"temperature": 0.05, "batch_size": 1},
+        {"temperature": 0.05},
         {"input": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), "target": torch.tensor([[1.0, 1.0, 0.0, 0.0]])},
         0.0,
     ],
@@ -60,12 +60,12 @@ def test_result(self, input_param, input_data, expected_val):
         np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
 
     def test_ill_shape(self):
-        loss = ContrastiveLoss(temperature=0.5, batch_size=1)
+        loss = ContrastiveLoss(temperature=0.5)
         with self.assertRaisesRegex(ValueError, ""):
             loss(torch.ones((1, 2, 3)), torch.ones((1, 1, 2, 3)))
 
     def test_with_cuda(self):
-        loss = ContrastiveLoss(temperature=0.5, batch_size=1)
+        loss = ContrastiveLoss(temperature=0.5)
         i = torch.ones((1, 10))
         j = torch.ones((1, 10))
         if torch.cuda.is_available():
@@ -74,6 +74,10 @@ def test_with_cuda(self):
         output = loss(i, j)
         np.testing.assert_allclose(output.detach().cpu().numpy(), 0.0, atol=1e-4, rtol=1e-4)
 
+    def check_warning_rasied(self):
+        with self.assertWarns(Warning):
+            ContrastiveLoss(temperature=0.5, batch_size=1)
+
 
 if __name__ == "__main__":
     unittest.main()