fix bugs with cifar100

Author: Demoren1

code/dependencies/darts_classification_module.py

@@ -9,12 +9,13 @@ def __init__(
         learning_rate: float = 0.001,
         weight_decay: float = 0.,
         auxiliary_loss_weight: float = 0.4,
-        max_epochs: int = 600
+        max_epochs: int = 600,
+        num_classes: int = 10
     ):
         self.auxiliary_loss_weight = auxiliary_loss_weight
         # Training length will be used in LR scheduler
         self.max_epochs = max_epochs
-        super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False, num_classes=10)
+        super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False, num_classes=num_classes)
 
     def configure_optimizers(self):
         """Customized optimizer with momentum, as well as a scheduler."""

code/output/ensemble_results.txt

@@ -1,5 +1,5 @@
-Ensemble Top-1 Accuracy: 10.01%
-Ensemble ECE: 0.0114
+Ensemble Top-1 Accuracy: 0.96%
+Ensemble ECE: 0.0015
 Number of models: 2
-Model 1 Accuracy: 9.90%
-Model 2 Accuracy: 10.02%
+Model 1 Accuracy: 1.00%
+Model 2 Accuracy: 0.90%

code/surrogate_hp.json

@@ -4,7 +4,7 @@
     "dataset_path": "third_dataset/",
     "device": "cpu",
     "developer_mode": true,
-    "n_models": 1300,
+    "n_models": 300,
 
     "upper_margin": 0.75,
     "lower_margin": 0.25,
@@ -36,14 +36,14 @@
     "n_models_in_pool": 128,
     "n_models_to_generate": 4096,
     "batch_size_inference": 4096,
-    "min_accuracy_for_pool": 0.85,
+    "min_accuracy_for_pool": 0.01,
     "plot_tsne": false,
     "best_models_save_path": "best_models/",
 
     "n_epochs_final": 1,
     "lr_final": 0.025,
     "batch_size_final": 96,
-    "dataset_name": "CIFAR10",
+    "dataset_name": "FashionMNIST",
     "final_dataset_path": "final_dataset/",
     "output_path": "output/",
     "width": 4,

code/surrogate_models/model_accuracy.pth

@@ -101,10 +101,13 @@ def get_data_loaders(self, batch_size: int = None) -> Tuple[DataLoader, DataLoad
 
         if self.config.dataset_name.lower() == "cifar10":
             dataset_cls = CIFAR10
+            self.num_classes = 10
         elif self.config.dataset_name.lower() == "cifar100":
             dataset_cls = CIFAR100
+            self.num_classes = 100
         elif self.config.dataset_name.lower() == "fashionmnist":
             dataset_cls = FashionMNIST
+            self.num_classes = 10
 
         train_data = nni.trace(dataset_cls)(
             root=self.config.final_dataset_path,
@@ -158,11 +161,20 @@ def train_model(
         Train a single model defined by architecture and return the trained model.
         """
         try:
+            if self.config.dataset_name.lower() == "cifar10":
+                dataset = "cifar"
+            elif self.config.dataset_name.lower() == "cifar100":
+                dataset = "cifar100"
+            elif self.config.dataset_name.lower() == "fashionmnist":
+                dataset = "cifar"
+            else:
+                raise ValueError(f"Unknown dataset: {self.config.dataset_name}")
+
             with model_context(architecture):
                 model = DartsSpace(
                     width=self.config.width,
                     num_cells=self.config.num_cells,
-                    dataset="cifar",
+                    dataset=dataset,
                 )
 
             model.to(self.device)
@@ -176,6 +188,7 @@ def train_model(
                     weight_decay=3e-4,
                     auxiliary_loss_weight=0.4,
                     max_epochs=self.config.n_epochs_final,
+                    num_classes=self.num_classes,
                 ),
                 trainer=Trainer(
                     gradient_clip_val=5.0,
@@ -289,7 +302,7 @@ def evaluate_ensemble(self, test_loader):
                 avg_output /= len(valid_models)
                 ensemble_probs = avg_output.softmax(dim=1)
                 confidences, preds_ens = ensemble_probs.max(1)
-                correct_ens_batch = (preds_ens == labels)
+                correct_ens_batch = preds_ens == labels
                 correct_ensemble += correct_ens_batch.sum().item()
 
                 confidences = confidences.cpu().float()