Upload spc.py

SPC-UQ/MNIST_Classification/trainers/spc.py

@@ -20,6 +20,9 @@ class SPC:
         self.model = model.to(device)
         self.criterion = nn.CrossEntropyLoss()
         self.optimizer = self._init_optimizer(optimizer_type, learning_rate)
+        self.optimizer_cls = self._init_optimizer_cls(optimizer_type, learning_rate)
+        self.optimizer_mar = self._init_optimizer_mar(optimizer_type, learning_rate)
+
 
         self.epoch = 0
         self.max_acc = -1
@@ -35,6 +38,22 @@ class SPC:
         else:
             raise ValueError(f"Unsupported optimizer type: {optimizer_type}")
 
+    def _init_optimizer_cls(self, optimizer_type, lr):
+        if optimizer_type.upper() == 'ADAM':
+            return optim.Adam(list(self.model.hidden.parameters()) + list(self.model.hidden_pred.parameters()) + list(self.model.output_pred.parameters()), lr=lr)
+        elif optimizer_type.upper() == 'SGD':
+            return optim.SGD(list(self.model.hidden.parameters()) + list(self.model.hidden_pred.parameters()) + list(self.model.output_pred.parameters()), lr=0.01, momentum=0.9, weight_decay=5e-4)
+        else:
+            raise ValueError(f"Unsupported optimizer type: {optimizer_type}")
+
+    def _init_optimizer_mar(self, optimizer_type, lr):
+        if optimizer_type.upper() == 'ADAM':
+            return optim.Adam(list(self.model.hidden_mar.parameters()) + list(self.model.output_mar.parameters()), lr=lr)
+        elif optimizer_type.upper() == 'SGD':
+            return optim.SGD(list(self.model.hidden_mar.parameters()) + list(self.model.output_mar.parameters()), lr=0.01, momentum=0.9, weight_decay=5e-4)
+        else:
+            raise ValueError(f"Unsupported optimizer type: {optimizer_type}")
+
     def cls_loss(self, target, logits):
         loss_cls = self.criterion(logits, target)
         return loss_cls
@@ -50,7 +69,7 @@ class SPC:
         loss_mar = F.mse_loss(mar, mar_target)
         return loss_mar
 
-    def run_train_step(self, data, target, num_classes):
+    def joint_train_step(self, data, target, num_classes):
         """Single training step."""
         self.model.train()
         logits, mar = self.model(data)
@@ -62,6 +81,28 @@ class SPC:
         self.optimizer.step()
         return loss.item(), mar
 
+    def cls_train_step(self, data, target):
+        """Single training step."""
+        self.model.train()
+        logits, mar = self.model(data)
+        loss_cls = self.cls_loss(target, logits)
+        loss = loss_cls
+        self.optimizer_cls.zero_grad()
+        loss.backward()
+        self.optimizer_cls.step()
+        return loss.item(), mar
+
+    def mar_train_step(self, data, target, num_classes):
+        """Single training step."""
+        self.model.train()
+        logits, mar = self.model(data)
+        loss_mar = self.mar_loss(target, num_classes, logits, mar)
+        loss = loss_mar
+        self.optimizer_mar.zero_grad()
+        loss.backward()
+        self.optimizer_mar.step()
+        return loss.item(), mar
+
     def evaluate(self, test_loader, ood_loader, num_classes, threshold=None):
         """
         Evaluate classification accuracy and uncertainty quantification:
@@ -73,6 +114,7 @@ class SPC:
         y_all, logits_all, mar_all = [], [], []
 
         with torch.no_grad():
+            start_time = time.time()
             for data, target in test_loader:
                 data, target = data.to(device), target.to(device)
                 logits, mar = self.model(data)
@@ -88,6 +130,7 @@ class SPC:
             expected_uncertainty = 2 * logits * (1 - logits)
             uncertainty = torch.sum(torch.abs(mar - expected_uncertainty), dim=1)
             confidences, predictions = torch.max(logits, dim=1)
+            test_time=(time.time() - start_time) * 1000
 
             threshold = np.quantile(uncertainty.cpu().numpy(), 0.5)
             acc = (predictions == y_all).float().mean().item()
@@ -119,10 +162,10 @@ class SPC:
             all_uncertainty = torch.cat([uncertainty, uncertainty_ood])
             auroc = metrics.roc_auc_score(bin_labels.cpu().numpy(), all_uncertainty.cpu().numpy())
 
-        return acc, acc_confident, acc_uncertain, auroc, 1000 * (len(test_loader.dataset) / len(test_loader))
+        return acc, acc_confident, acc_uncertain, auroc, test_time / len(test_loader)
 
     def train(self, train_dataset, test_dataset, ood_dataset, num_classes,
-              batch_size=128, num_epochs=40, verbose=True, freq=1):
+              batch_size=128, num_epochs=40, verbose=True, freq=1, joint_training=0):
         """
         Train classifier with SPC-based uncertainty, and evaluate per epoch.
         """
@@ -134,38 +177,84 @@ class SPC:
         acc_curve, acc_conf_curve, acc_unc_curve = [], [], []
         train_times, test_times = [], []
 
-        for self.epoch in range(1, num_epochs + 1):
-            total_loss, total_unc = 0.0, 0.0
-            start_time = time.time()
+        if joint_training:
+            for self.epoch in range(1, num_epochs + 1):
+                total_loss, total_unc = 0.0, 0.0
+                start_time = time.time()
+
+                for data, target in train_loader:
+                    data, target = data.to(device), target.to(device)
+                    loss, mar = self.joint_train_step(data, target, num_classes)
+                    total_loss += loss
+
+                    if verbose:
+                        with torch.no_grad():
+                            logits, mar = self.model(data)
+                            prob = F.softmax(logits, dim=1)
+                            mar_target = 2 * prob * (1 - prob)
+                            uncertainty = torch.sum(torch.abs(mar - mar_target), dim=1)
+                            total_unc += uncertainty.mean().item()
+
+                train_times.append((time.time() - start_time) * 1000)
+                avg_loss = total_loss / len(train_loader)
+                avg_uncertainty = total_unc / len(train_loader)
+
+                loss_curve.append(avg_loss)
+                uncertainty_curve.append(avg_uncertainty)
+
+                if self.epoch % freq == 0:
+                    acc, acc_conf, acc_unc, auroc, test_time = self.evaluate(test_loader, ood_loader, num_classes)
+                    acc_curve.append(acc)
+                    acc_conf_curve.append(acc_conf)
+                    acc_unc_curve.append(acc_unc)
+                    test_times.append(test_time)
+
+                    if acc > self.max_acc:
+                        self.max_acc = acc
+                        self.max_acc_confident = acc_conf
+                        self.max_acc_uncertain = acc_unc
+                        self.max_auroc = auroc
+
+        else:
+            for self.epoch in range(1, num_epochs + 1):
+                total_loss = 0.0
+
+                for data, target in train_loader:
+                    data, target = data.to(device), target.to(device)
+                    loss, mar = self.cls_train_step(data, target)
+                    total_loss += loss
+
+                avg_loss = total_loss / len(train_loader)
+                loss_curve.append(avg_loss)
+
+            for self.epoch in range(1, num_epochs + 1):
+                total_unc = 0.0
+                start_time = time.time()
+
+                for data, target in train_loader:
+                    data, target = data.to(device), target.to(device)
+                    loss, mar = self.mar_train_step(data, target, num_classes)
+
+                    if verbose:
+                        with torch.no_grad():
+                            logits, mar = self.model(data)
+                            prob = F.softmax(logits, dim=1)
+                            mar_target = 2 * prob * (1 - prob)
+                            uncertainty = torch.sum(torch.abs(mar - mar_target), dim=1)
+                            total_unc += uncertainty.mean().item()
+
+                train_times.append((time.time() - start_time) * 1000)
+                avg_uncertainty = total_unc / len(train_loader)
+
+                uncertainty_curve.append(avg_uncertainty)
+
+                if self.epoch % freq == 0:
+                    acc, acc_conf, acc_unc, auroc, test_time = self.evaluate(test_loader, ood_loader, num_classes)
+                    acc_curve.append(acc)
+                    acc_conf_curve.append(acc_conf)
+                    acc_unc_curve.append(acc_unc)
+                    test_times.append(test_time)
 
-            for data, target in train_loader:
-                data, target = data.to(device), target.to(device)
-                loss, mar = self.run_train_step(data, target, num_classes)
-                total_loss += loss
-
-                if verbose:
-                    with torch.no_grad():
-                        logits, mar = self.model(data)
-                        prob = F.softmax(logits, dim=1)
-                        mar_target = 2 * prob * (1 - prob)
-                        uncertainty = torch.sum(torch.abs(mar - mar_target), dim=1)
-                        total_unc += uncertainty.mean().item()
-
-            train_times.append((time.time() - start_time) * 1000)
-            avg_loss = total_loss / len(train_loader)
-            avg_uncertainty = total_unc / len(train_loader)
-
-            loss_curve.append(avg_loss)
-            uncertainty_curve.append(avg_uncertainty)
-
-            if self.epoch % freq == 0:
-                acc, acc_conf, acc_unc, auroc, test_time = self.evaluate(test_loader, ood_loader, num_classes)
-                acc_curve.append(acc)
-                acc_conf_curve.append(acc_conf)
-                acc_unc_curve.append(acc_unc)
-                test_times.append(test_time)
-
-                if acc > self.max_acc:
                     self.max_acc = acc
                     self.max_acc_confident = acc_conf
                     self.max_acc_uncertain = acc_unc
@@ -193,7 +282,7 @@ class SPC:
         plt.grid(True); plt.legend(); plt.show()
 
         plt.figure(figsize=(10, 6))
-        plt.plot(acc, label='Test Accuracy')
+        plt.plot(acc, label='Overall Accuracy')
         plt.plot(acc_cer, label='Confident Accuracy')
         plt.plot(acc_unc, label='Uncertain Accuracy')
         plt.title('Test Accuracy Curves')