diff options
| -rw-r--r-- | libs/criteria.py | 10 | ||||
| -rw-r--r-- | simclr/main.py | 18 |
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() |