Updated Contrastive Loss batch Size

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,22 +68,23 @@ 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)
         denominator = negatives_mask * torch.exp(sim_matrix / temperature_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)

tests/test_contrastive_loss.py

@@ -19,33 +19,33 @@
 
 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]]),
         },
         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]]),
         },
         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()