Optimize batching

2 files changed, 13 insertions, 15 deletions

diff --git a/libs/criteria.py b/libs/criteria.py
index 6954cf3..baa36ce 100644
--- a/libs/criteria.py
+++ b/libs/criteria.py
@@ -8,17 +8,17 @@ class InfoNCELoss(nn.Module):
         super().__init__()
         self.temp = temp
 
-    def forward(self, feat1: Tensor, feat2: Tensor) -> tuple[Tensor, Tensor]:
-        bz = feat1.size(0)
-        feat1_norm = F.normalize(feat1)
-        feat2_norm = F.normalize(feat2)
+    def forward(self, feature: Tensor) -> tuple[Tensor, Tensor]:
+        bz = feature.size(0) // 2
+        feat_norm = F.normalize(feature)
+        feat1_norm, feat2_norm = feat_norm.split(bz)
         logits = feat1_norm @ feat2_norm.T
         pos_logits_mask = torch.eye(bz, dtype=torch.bool)
         pos_logits = logits[pos_logits_mask].unsqueeze(-1)
         neg_logits = logits[~pos_logits_mask].view(bz, -1)
         # Put the positive at first (0-th) and maximize its likelihood
         logits = torch.cat([pos_logits, neg_logits], dim=1)
-        labels = torch.zeros(bz, dtype=torch.long, device=feat1.device)
+        labels = torch.zeros(bz, dtype=torch.long, device=feature.device)
         loss_contra = F.cross_entropy(logits / self.temp, labels)
         acc_contra = (logits.argmax(dim=1) == labels).float().mean()
 
diff --git a/simclr/main.py b/simclr/main.py
index 93a2d83..b170355 100644
--- a/simclr/main.py
+++ b/simclr/main.py
@@ -294,12 +294,11 @@ class SimCLRTrainer(Trainer):
         train_loader = iter(self.train_loader)
         model.train()
         for iter_ in range(self.restore_iter, num_iters):
-            (input1, input2), _ = next(train_loader)
-            input1, input2 = input1.to(device), input2.to(device)
+            input_, _ = next(train_loader)
+            input_ = torch.cat(input_).to(device)
             model.zero_grad()
-            output1 = model(input1)
-            output2 = model(input2)
-            train_loss, train_accuracy = loss_fn(output1, output2)
+            output = model(input_)
+            train_loss, train_accuracy = loss_fn(output)
             train_loss.backward()
             optim.step()
             self.log(logger, self.BatchLogRecord(
@@ -327,11 +326,10 @@ class SimCLRTrainer(Trainer):
         model = self.models['model']
         model.eval()
         with torch.no_grad():
-            for (input1, input2), _ in self.test_loader:
-                input1, input2 = input1.to(device), input2.to(device)
-                output1 = model(input1)
-                output2 = model(input2)
-                loss, accuracy = loss_fn(output1, output2)
+            for input_, _ in self.test_loader:
+                input_ = torch.cat(input_).to(device)
+                output = model(input_)
+                loss, accuracy = loss_fn(output)
                 yield loss.item(), accuracy.item()