ocsvm working

Deep-SAD-PyTorch/src/DeepSAD.py

@@ -27,9 +27,11 @@ class DeepSAD(object):
         ae_results: A dictionary to save the autoencoder results.
     """
 
-    def __init__(self, eta: float = 1.0):
+    def __init__(self, rep_dim: int, eta: float = 1.0):
         """Inits DeepSAD with hyperparameter eta."""
 
+        self.rep_dim = rep_dim  # representation dimension
+
         self.eta = eta
         self.c = None  # hypersphere center c
 
@@ -89,10 +91,10 @@ class DeepSAD(object):
 
         self.ae_results = {"train_time": None, "test_auc": None, "test_time": None}
 
-    def set_network(self, net_name, rep_dim=1024):
+    def set_network(self, net_name):
         """Builds the neural network phi."""
         self.net_name = net_name
-        self.net = build_network(net_name, rep_dim=rep_dim)
+        self.net = build_network(net_name, self.rep_dim)
 
     def train(
         self,
@@ -240,7 +242,7 @@ class DeepSAD(object):
         """Pretrains the weights for the Deep SAD network phi via autoencoder."""
 
         # Set autoencoder network
-        self.ae_net = build_autoencoder(self.net_name)
+        self.ae_net = build_autoencoder(self.net_name, self.rep_dim)
 
         # Train
         self.ae_optimizer_name = optimizer_name
@@ -340,7 +342,7 @@ class DeepSAD(object):
             # json.dump(self.results, fp)
             pickle.dump(self.results, fp)
 
-    def save_ae_results(self, export_json):
+    def save_ae_results(self, export_pkl):
         """Save autoencoder results dict to a JSON-file."""
-        with open(export_json, "w") as fp:
+        with open(export_pkl, "wb") as fp:
-            json.dump(self.ae_results, fp)
+            pickle.dump(self.ae_results, fp)

Deep-SAD-PyTorch/src/base/base_net.py

@@ -35,3 +35,5 @@ class BaseNet(nn.Module):
         self.logger.info(
             torchscan.summary(self, self.input_dim, receptive_field=receptive_field)
         )
+        module_info = torchscan.crawl_module(self, self.input_dim)
+        pass

Deep-SAD-PyTorch/src/baselines/ocsvm.py

@@ -1,4 +1,3 @@
-import json
 import logging
 import pickle
 import time
@@ -20,7 +19,7 @@ from networks.main import build_autoencoder
 class OCSVM(object):
     """A class for One-Class SVM models."""
 
-    def __init__(self, kernel="rbf", nu=0.1, hybrid=False):
+    def __init__(self, kernel="rbf", nu=0.1, hybrid=False, latent_space_dim=128):
         """Init OCSVM instance."""
         self.kernel = kernel
         self.nu = nu
@@ -30,6 +29,7 @@ class OCSVM(object):
         self.model = OneClassSVM(kernel=kernel, nu=nu, verbose=True, max_mem_size=4048)
 
         self.hybrid = hybrid
+        self.latent_space_dim = latent_space_dim
         self.ae_net = None  # autoencoder network for the case of a hybrid model
         self.linear_model = (
             None  # also init a model with linear kernel if hybrid approach
@@ -38,8 +38,16 @@ class OCSVM(object):
         self.results = {
             "train_time": None,
             "test_time": None,
-            "test_auc": None,
+            "test_auc_exp_based": None,
-            "test_scores": None,
+            "test_roc_exp_based": None,
+            "test_prc_exp_based": None,
+            "test_ap_exp_based": None,
+            "test_scores_exp_based": None,
+            "test_auc_manual_based": None,
+            "test_roc_manual_based": None,
+            "test_prc_manual_based": None,
+            "test_ap_manual_based": None,
+            "test_scores_manual_based": None,
             "train_time_linear": None,
             "test_time_linear": None,
             "test_auc_linear": None,
@@ -70,15 +78,11 @@ class OCSVM(object):
         # Get data from loader
         X = ()
         for data in train_loader:
-            inputs, _, _, _, _ = data
+            inputs, _, _, _, _, _ = data  # Updated unpacking
             inputs = inputs.to(device)
             if self.hybrid:
-                inputs = self.ae_net.encoder(
+                inputs = self.ae_net.encoder(inputs)
-                    inputs
+            X_batch = inputs.view(inputs.size(0), -1)
-                )  # in hybrid approach, take code representation of AE as features
-            X_batch = inputs.view(
-                inputs.size(0), -1
-            )  # X_batch.shape = (batch_size, n_channels * height * width)
             X += (X_batch.cpu().data.numpy(),)
         X = np.concatenate(X)
 
@@ -101,40 +105,59 @@ class OCSVM(object):
                 batch_size=batch_size, num_workers=n_jobs_dataloader
             )
 
+        # Sample hold-out set from test set
         X_test = ()
-        labels = []
+        labels_exp = []
+        labels_manual = []
+
         for data in test_loader:
-            inputs, label_batch, _, _, _ = data
+            inputs, label_exp, label_manual, _, _, _ = data  # Updated unpacking
-            inputs, label_batch = inputs.to(device), label_batch.to(device)
+            inputs = inputs.to(device)
+            label_exp = label_exp.to(device)
+            label_manual = label_manual.to(device)
+
             if self.hybrid:
-                inputs = self.ae_net.encoder(
+                inputs = self.ae_net.encoder(inputs)
-                    inputs
+            X_batch = inputs.view(inputs.size(0), -1)
-                )  # in hybrid approach, take code representation of AE as features
-            X_batch = inputs.view(
-                inputs.size(0), -1
-            )  # X_batch.shape = (batch_size, n_channels * height * width)
             X_test += (X_batch.cpu().data.numpy(),)
-            labels += label_batch.cpu().data.numpy().astype(np.int64).tolist()
+            labels_exp += label_exp.cpu().data.numpy().astype(np.int64).tolist()
-        X_test, labels = np.concatenate(X_test), np.array(labels)
+            labels_manual += label_manual.cpu().data.numpy().astype(np.int64).tolist()
-        n_test, n_normal, n_outlier = (
+
-            len(X_test),
+        X_test = np.concatenate(X_test)
-            np.sum(labels == 0),
+        labels_exp = np.array(labels_exp)
-            np.sum(labels == 1),
+        labels_manual = np.array(labels_manual)
-        )
+
-        n_val = int(0.1 * n_test)
+        # Use experiment-based labels for model selection (could also use manual-based)
-        n_val_normal, n_val_outlier = (
+        labels = labels_exp
-            int(n_val * (n_normal / n_test)),
+
-            int(n_val * (n_outlier / n_test)),
+        # Count samples for validation split (-1: anomaly, 1: normal, 0: unknown)
-        )
+        n_test = len(X_test)
-        perm = np.random.permutation(n_test)
+        n_normal = np.sum(labels == 1)
+        n_outlier = np.sum(labels == -1)
+        n_val = int(0.1 * n_test)  # Use 10% of test data for validation
+
+        # Only consider labeled samples for validation
+        valid_mask = labels != 0
+        X_test_valid = X_test[valid_mask]
+        labels_valid = labels[valid_mask]
+
+        # Calculate validation split sizes
+        n_val_normal = int(n_val * (n_normal / (n_normal + n_outlier)))
+        n_val_outlier = int(n_val * (n_outlier / (n_normal + n_outlier)))
+
+        # Create validation set with balanced normal/anomaly ratio
+        perm = np.random.permutation(len(X_test_valid))
         X_val = np.concatenate(
             (
-                X_test[perm][labels[perm] == 0][:n_val_normal],
+                X_test_valid[perm][labels_valid[perm] == 1][:n_val_normal],
-                X_test[perm][labels[perm] == 1][:n_val_outlier],
+                X_test_valid[perm][labels_valid[perm] == -1][:n_val_outlier],
             )
         )
-        labels = np.array([0] * n_val_normal + [1] * n_val_outlier)
+        val_labels = np.array(
+            [0] * n_val_normal + [1] * n_val_outlier
+        )  # Convert to binary (0: normal, 1: anomaly)
 
+        # Model selection loop
         i = 1
         for gamma in gammas:
             # Model candidate
@@ -155,12 +178,12 @@ class OCSVM(object):
             scores = (-1.0) * model.decision_function(X_val)
             scores = scores.flatten()
 
-            # Compute AUC
+            # Compute AUC with binary labels
-            auc = roc_auc_score(labels, scores)
+            auc = roc_auc_score(val_labels, scores)
 
             logger.info(
                 f"  | Model {i:02}/{len(gammas):02} | Gamma: {gamma:.8f} | Train Time: {train_time:.3f}s "
-                f"| Val AUC: {100. * auc:.2f} |"
+                f"| Val AUC: {100.0 * auc:.2f} |"
             )
 
             if auc > best_auc:
@@ -182,7 +205,7 @@ class OCSVM(object):
             self.results["train_time_linear"] = train_time
 
         logger.info(
-            f"Best Model: | Gamma: {self.gamma:.8f} | AUC: {100. * best_auc:.2f}"
+            f"Best Model: | Gamma: {self.gamma:.8f} | AUC: {100.0 * best_auc:.2f}"
         )
         logger.info("Training Time: {:.3f}s".format(self.results["train_time"]))
         logger.info("Finished training.")
@@ -210,51 +233,121 @@ class OCSVM(object):
             )
 
         # Get data from loader
-        idx_label_score = []
+        idx_label_score_exp = []
+        idx_label_score_manual = []
         X = ()
         idxs = []
-        labels = []
+        labels_exp = []
+        labels_manual = []
+
         for data in test_loader:
-            inputs, label_batch, _, idx, _ = data
+            inputs, label_exp, label_manual, _, idx, _ = data  # Updated unpacking
-            inputs, label_batch, idx = (
+            inputs, label_exp, label_manual, idx = (
                 inputs.to(device),
-                label_batch.to(device),
+                label_exp.to(device),
+                label_manual.to(device),
                 idx.to(device),
             )
+
             if self.hybrid:
-                inputs = self.ae_net.encoder(
+                inputs = self.ae_net.encoder(inputs)
-                    inputs
+            X_batch = inputs.view(inputs.size(0), -1)
-                )  # in hybrid approach, take code representation of AE as features
-            X_batch = inputs.view(
-                inputs.size(0), -1
-            )  # X_batch.shape = (batch_size, n_channels * height * width)
             X += (X_batch.cpu().data.numpy(),)
             idxs += idx.cpu().data.numpy().astype(np.int64).tolist()
-            labels += label_batch.cpu().data.numpy().astype(np.int64).tolist()
+            labels_exp += label_exp.cpu().data.numpy().astype(np.int64).tolist()
+            labels_manual += label_manual.cpu().data.numpy().astype(np.int64).tolist()
+
         X = np.concatenate(X)
+        labels_exp = np.array(labels_exp)
+        labels_manual = np.array(labels_manual)
+
+        # Count and log label stats for exp_based
+        n_exp_normal = np.sum(labels_exp == 1)
+        n_exp_anomaly = np.sum(labels_exp == -1)
+        n_exp_unknown = np.sum(labels_exp == 0)
+        logger.info(
+            f"Exp-based labels: normal(1)={n_exp_normal}, "
+            f"anomaly(-1)={n_exp_anomaly}, unknown(0)={n_exp_unknown}"
+        )
+
+        # Count and log label stats for manual_based
+        n_manual_normal = np.sum(labels_manual == 1)
+        n_manual_anomaly = np.sum(labels_manual == -1)
+        n_manual_unknown = np.sum(labels_manual == 0)
+        logger.info(
+            f"Manual-based labels: normal(1)={n_manual_normal}, "
+            f"anomaly(-1)={n_manual_anomaly}, unknown(0)={n_manual_unknown}"
+        )
 
         # Testing
         logger.info("Starting testing...")
         start_time = time.time()
-
         scores = (-1.0) * self.model.decision_function(X)
-
         self.results["test_time"] = time.time() - start_time
         scores = scores.flatten()
-        self.rho = -self.model.intercept_[0]
 
-        # Save triples of (idx, label, score) in a list
+        # Save triples of (idx, label, score) for both label types
-        idx_label_score += list(zip(idxs, labels, scores.tolist()))
+        idx_label_score_exp += list(zip(idxs, labels_exp.tolist(), scores.tolist()))
-        self.results["test_scores"] = idx_label_score
+        idx_label_score_manual += list(
+            zip(idxs, labels_manual.tolist(), scores.tolist())
+        )
 
-        # Compute AUC
+        self.results["test_scores_exp_based"] = idx_label_score_exp
-        _, labels, scores = zip(*idx_label_score)
+        self.results["test_scores_manual_based"] = idx_label_score_manual
-        labels = np.array(labels)
+
-        scores = np.array(scores)
+        # --- Evaluation for exp_based (only labeled samples) ---
-        self.results["test_auc"] = roc_auc_score(labels, scores)
+        valid_mask_exp = labels_exp != 0
-        self.results["test_roc"] = roc_curve(labels, scores)
+        if np.any(valid_mask_exp):
-        self.results["test_prc"] = precision_recall_curve(labels, scores)
+            labels_exp_binary = (labels_exp[valid_mask_exp] == -1).astype(int)
-        self.results["test_ap"] = average_precision_score(labels, scores)
+            scores_exp_valid = scores[valid_mask_exp]
+
+            self.results["test_auc_exp_based"] = roc_auc_score(
+                labels_exp_binary, scores_exp_valid
+            )
+            self.results["test_roc_exp_based"] = roc_curve(
+                labels_exp_binary, scores_exp_valid
+            )
+            self.results["test_prc_exp_based"] = precision_recall_curve(
+                labels_exp_binary, scores_exp_valid
+            )
+            self.results["test_ap_exp_based"] = average_precision_score(
+                labels_exp_binary, scores_exp_valid
+            )
+
+            logger.info(
+                "Test AUC (exp_based): {:.2f}%".format(
+                    100.0 * self.results["test_auc_exp_based"]
+                )
+            )
+        else:
+            logger.info("Test AUC (exp_based): N/A (no labeled samples)")
+
+        # --- Evaluation for manual_based (only labeled samples) ---
+        valid_mask_manual = labels_manual != 0
+        if np.any(valid_mask_manual):
+            labels_manual_binary = (labels_manual[valid_mask_manual] == -1).astype(int)
+            scores_manual_valid = scores[valid_mask_manual]
+
+            self.results["test_auc_manual_based"] = roc_auc_score(
+                labels_manual_binary, scores_manual_valid
+            )
+            self.results["test_roc_manual_based"] = roc_curve(
+                labels_manual_binary, scores_manual_valid
+            )
+            self.results["test_prc_manual_based"] = precision_recall_curve(
+                labels_manual_binary, scores_manual_valid
+            )
+            self.results["test_ap_manual_based"] = average_precision_score(
+                labels_manual_binary, scores_manual_valid
+            )
+
+            logger.info(
+                "Test AUC (manual_based): {:.2f}%".format(
+                    100.0 * self.results["test_auc_manual_based"]
+                )
+            )
+        else:
+            logger.info("Test AUC (manual_based): N/A (no labeled samples)")
 
         # If hybrid, also test model with linear kernel
         if self.hybrid:
@@ -262,35 +355,115 @@ class OCSVM(object):
             scores_linear = (-1.0) * self.linear_model.decision_function(X)
             self.results["test_time_linear"] = time.time() - start_time
             scores_linear = scores_linear.flatten()
-            self.results["test_auc_linear"] = roc_auc_score(labels, scores_linear)
+            # Save linear model results for both label types
+            # --- exp_based ---
+            valid_mask_exp_linear = labels_exp != 0
+            if np.any(valid_mask_exp_linear):
+                labels_exp_binary_linear = (
+                    labels_exp[valid_mask_exp_linear] == -1
+                ).astype(int)
+                scores_exp_linear_valid = scores_linear[valid_mask_exp_linear]
+                self.results["test_auc_linear_exp_based"] = roc_auc_score(
+                    labels_exp_binary_linear, scores_exp_linear_valid
+                )
+                self.results["test_roc_linear_exp_based"] = roc_curve(
+                    labels_exp_binary_linear, scores_exp_linear_valid
+                )
+                self.results["test_prc_linear_exp_based"] = precision_recall_curve(
+                    labels_exp_binary_linear, scores_exp_linear_valid
+                )
+                self.results["test_ap_linear_exp_based"] = average_precision_score(
+                    labels_exp_binary_linear, scores_exp_linear_valid
+                )
+            else:
+                self.results["test_auc_linear_exp_based"] = None
+                self.results["test_roc_linear_exp_based"] = None
+                self.results["test_prc_linear_exp_based"] = None
+                self.results["test_ap_linear_exp_based"] = None
+
+            # --- manual_based ---
+            valid_mask_manual_linear = labels_manual != 0
+            if np.any(valid_mask_manual_linear):
+                labels_manual_binary_linear = (
+                    labels_manual[valid_mask_manual_linear] == -1
+                ).astype(int)
+                scores_manual_linear_valid = scores_linear[valid_mask_manual_linear]
+                self.results["test_auc_linear_manual_based"] = roc_auc_score(
+                    labels_manual_binary_linear, scores_manual_linear_valid
+                )
+                self.results["test_roc_linear_manual_based"] = roc_curve(
+                    labels_manual_binary_linear, scores_manual_linear_valid
+                )
+                self.results["test_prc_linear_manual_based"] = precision_recall_curve(
+                    labels_manual_binary_linear, scores_manual_linear_valid
+                )
+                self.results["test_ap_linear_manual_based"] = average_precision_score(
+                    labels_manual_binary_linear, scores_manual_linear_valid
+                )
+            else:
+                self.results["test_auc_linear_manual_based"] = None
+                self.results["test_roc_linear_manual_based"] = None
+                self.results["test_prc_linear_manual_based"] = None
+                self.results["test_ap_linear_manual_based"] = None
+            # Log exp_based results for linear model
+            if self.results["test_auc_linear_exp_based"] is not None:
                 logger.info(
-                "Test AUC linear model: {:.2f}%".format(
+                    "Test AUC linear model (exp_based): {:.2f}%".format(
-                    100.0 * self.results["test_auc_linear"]
+                        100.0 * self.results["test_auc_linear_exp_based"]
                     )
                 )
+            else:
+                logger.info(
+                    "Test AUC linear model (exp_based): N/A (no labeled samples)"
+                )
+
+            # Log manual_based results for linear model
+            if self.results["test_auc_linear_manual_based"] is not None:
+                logger.info(
+                    "Test AUC linear model (manual_based): {:.2f}%".format(
+                        100.0 * self.results["test_auc_linear_manual_based"]
+                    )
+                )
+            else:
+                logger.info(
+                    "Test AUC linear model (manual_based): N/A (no labeled samples)"
+                )
+
             logger.info(
                 "Test Time linear model: {:.3f}s".format(
                     self.results["test_time_linear"]
                 )
             )
 
-        # Log results
+        # Log results for both label types
-        logger.info("Test AUC: {:.2f}%".format(100.0 * self.results["test_auc"]))
+        if self.results.get("test_auc_exp_based") is not None:
+            logger.info(
+                "Test AUC (exp_based): {:.2f}%".format(
+                    100.0 * self.results["test_auc_exp_based"]
+                )
+            )
+        else:
+            logger.info("Test AUC (exp_based): N/A (no labeled samples)")
+
+        if self.results.get("test_auc_manual_based") is not None:
+            logger.info(
+                "Test AUC (manual_based): {:.2f}%".format(
+                    100.0 * self.results["test_auc_manual_based"]
+                )
+            )
+        else:
+            logger.info("Test AUC (manual_based): N/A (no labeled samples)")
         logger.info("Test Time: {:.3f}s".format(self.results["test_time"]))
         logger.info("Finished testing.")
 
-    def load_ae(self, dataset_name, model_path):
+    def load_ae(self, model_path, net_name, device="cpu"):
         """Load pretrained autoencoder from model_path for feature extraction in a hybrid OC-SVM model."""
 
         model_dict = torch.load(model_path, map_location="cpu")
         ae_net_dict = model_dict["ae_net_dict"]
-        if dataset_name in ["mnist", "fmnist", "cifar10"]:
-            net_name = dataset_name + "_LeNet"
-        else:
-            net_name = dataset_name + "_mlp"
 
         if self.ae_net is None:
-            self.ae_net = build_autoencoder(net_name)
+            self.ae_net = build_autoencoder(net_name, rep_dim=self.latent_space_dim)
 
         # update keys (since there was a change in network definition)
         ae_keys = list(self.ae_net.state_dict().keys())
@@ -301,6 +474,8 @@ class OCSVM(object):
             i += 1
 
         self.ae_net.load_state_dict(ae_net_dict)
+        if device != "cpu":
+            self.ae_net.to(torch.device(device))
         self.ae_net.eval()
 
     def save_model(self, export_path):

Deep-SAD-PyTorch/src/main.py

@@ -167,6 +167,12 @@ from utils.visualization.plot_images_grid import plot_images_grid
     default=1e-6,
     help="Weight decay (L2 penalty) hyperparameter for Deep SAD objective.",
 )
+@click.option(
+    "--latent_space_dim",
+    type=int,
+    default=128,
+    help="Dimensionality of the latent space for the autoencoder.",
+)
 @click.option(
     "--pretrain",
     type=bool,
@@ -303,6 +309,7 @@ def main(
     lr_milestone,
     batch_size,
     weight_decay,
+    latent_space_dim,
     pretrain,
     ae_optimizer_name,
     ae_lr,
@@ -415,7 +422,7 @@ def main(
         train_passes = range(k_fold_num) if k_fold else [None]
 
         train_isoforest = True
-        train_ocsvm = False
+        train_ocsvm = True
         train_deepsad = True
 
         for fold_idx in train_passes:
@@ -424,10 +431,6 @@ def main(
             else:
                 logger.info(f"Fold {fold_idx + 1}/{k_fold_num}")
 
-            # Initialize OC-SVM model
-            if train_ocsvm:
-                ocsvm = OCSVM(kernel=ocsvm_kernel, nu=ocsvm_nu, hybrid=False)
-
             # Initialize Isolation Forest model
             if train_isoforest:
                 Isoforest = IsoForest(
@@ -441,7 +444,7 @@ def main(
 
             # Initialize DeepSAD model and set neural network phi
             if train_deepsad:
-                deepSAD = DeepSAD(cfg.settings["eta"])
+                deepSAD = DeepSAD(latent_space_dim, cfg.settings["eta"])
                 deepSAD.set_network(net_name)
 
             # If specified, load Deep SAD model (center c, network weights, and possibly autoencoder weights)
@@ -486,10 +489,31 @@ def main(
 
                 # Save pretraining results
                 if fold_idx is None:
-                    deepSAD.save_ae_results(export_json=xp_path + "/ae_results.json")
+                    deepSAD.save_ae_results(export_pkl=xp_path + "/ae_results.pkl")
+                    ae_model_path = xp_path + "/ae_model.tar"
+                    deepSAD.save_model(export_model=ae_model_path, save_ae=True)
                 else:
                     deepSAD.save_ae_results(
-                        export_json=xp_path + f"/ae_results_{fold_idx}.json"
+                        export_pkl=xp_path + f"/ae_results_{fold_idx}.pkl"
+                    )
+                    ae_model_path = xp_path + f"/ae_model_{fold_idx}.tar"
+                    deepSAD.save_model(export_model=ae_model_path, save_ae=True)
+
+            # Initialize OC-SVM model (after pretraining to use autoencoder features)
+            if train_ocsvm:
+                ocsvm = OCSVM(
+                    kernel=ocsvm_kernel,
+                    nu=ocsvm_nu,
+                    hybrid=True,
+                    latent_space_dim=latent_space_dim,
+                )
+                if load_model and not pretrain:
+                    ae_model_path = load_model
+                ocsvm.load_ae(
+                    net_name=net_name, model_path=ae_model_path, device=device
+                )
+                logger.info(
+                    f"Loaded pretrained autoencoder for features from {ae_model_path}."
                 )
 
             # Log training details
@@ -525,7 +549,7 @@ def main(
                     device=device,
                     n_jobs_dataloader=n_jobs_dataloader,
                     k_fold_idx=fold_idx,
-                    batch_size=8,
+                    batch_size=256,
                 )
 
             # Train model on dataset
@@ -553,7 +577,7 @@ def main(
                     device=device,
                     n_jobs_dataloader=n_jobs_dataloader,
                     k_fold_idx=fold_idx,
-                    batch_size=8,
+                    batch_size=256,
                 )
 
             # Test model
@@ -730,7 +754,7 @@ def main(
             logger.info("Known anomaly classes: %s" % (dataset.known_outlier_classes,))
 
         # Initialize DeepSAD model and set neural network phi
-        deepSAD = DeepSAD(cfg.settings["eta"])
+        deepSAD = DeepSAD(latent_space_dim, cfg.settings["eta"])
         deepSAD.set_network(net_name)
 
         # If specified, load Deep SAD model (center c, network weights, and possibly autoencoder weights)
@@ -776,22 +800,39 @@ def main(
             ratio_known_outlier,
             ratio_pollution,
             random_state=np.random.RandomState(cfg.settings["seed"]),
+            k_fold_num=k_fold_num,
         )
 
-        # Dictionary to store results for each dimension
+        # Set up k-fold passes
-        # ae_elbow_dims = [32, 64, 128, 256, 384, 512, 768, 1024]
+        train_passes = range(k_fold_num) if k_fold else [None]
-        ae_elbow_dims = [32, 64]
+
-        elbow_results = {"dimensions": list(ae_elbow_dims), "ae_results": {}}
+        # Test dimensions
+        ae_elbow_dims = [32, 64, 128, 256, 384, 512, 768, 1024]
 
         # Test each dimension
         for rep_dim in ae_elbow_dims:
             logger.info(f"Testing autoencoder with latent dimension: {rep_dim}")
 
+            # Results dictionary for this dimension
+            dim_results = {
+                "dimension": rep_dim,
+                "ae_results": {},
+                "k_fold": k_fold,
+                "k_fold_num": k_fold_num,
+            }
+
+            # For each fold
+            for fold_idx in train_passes:
+                if fold_idx is None:
+                    logger.info(f"Dimension {rep_dim}: Single training without k-fold")
+                else:
+                    logger.info(
+                        f"Dimension {rep_dim}: Fold {fold_idx + 1}/{k_fold_num}"
+                    )
+
                 # Initialize DeepSAD model with current dimension
-            deepSAD = DeepSAD(cfg.settings["eta"])
+                deepSAD = DeepSAD(rep_dim, cfg.settings["eta"])
-            deepSAD.set_network(
+                deepSAD.set_network(net_name)
-                net_name, rep_dim=rep_dim
-            )  # Pass rep_dim to network builder
 
                 # Pretrain autoencoder with current dimension
                 deepSAD.pretrain(
@@ -804,24 +845,40 @@ def main(
                     weight_decay=cfg.settings["ae_weight_decay"],
                     device=device,
                     n_jobs_dataloader=n_jobs_dataloader,
+                    k_fold_idx=fold_idx,
                 )
 
-            # Store results for this dimension
+                # Store results for this fold
-            elbow_results["ae_results"][rep_dim] = deepSAD.ae_results
+                fold_key = "single" if fold_idx is None else f"fold_{fold_idx}"
+                dim_results["ae_results"][fold_key] = deepSAD.ae_results
 
-            logger.info(f"Finished testing dimension {rep_dim}")
+                logger.info(
+                    f"Finished testing dimension {rep_dim} "
+                    + (
+                        f"fold {fold_idx + 1}/{k_fold_num}"
+                        if fold_idx is not None
+                        else "single pass"
+                    )
+                )
 
                 # Clear some memory
                 del deepSAD
                 torch.cuda.empty_cache()
 
-        # Save all results
+            # Save results for this dimension (includes all folds)
-        results_path = Path(xp_path) / f"ae_elbow_results_{net_name}.pkl"
+            results_filename = (
+                f"ae_elbow_results_{net_name}_dim_{rep_dim}"
+                + ("_kfold" if k_fold else "")
+                + ".pkl"
+            )
+            results_path = Path(xp_path) / results_filename
+
             with open(results_path, "wb") as f:
-            pickle.dump(elbow_results, f)
+                pickle.dump(dim_results, f)
-
-        logger.info(f"Saved elbow test results to {results_path}")
 
+            logger.info(
+                f"Saved elbow test results for dimension {rep_dim} to {results_path}"
+            )
         else:
             logger.error(f"Unknown action: {action}")
 

Deep-SAD-PyTorch/src/networks/main.py

@@ -9,7 +9,7 @@ from .subter_LeNet_Split import SubTer_LeNet_Split, SubTer_LeNet_Split_Autoencod
 from .vae import VariationalAutoencoder
 
 
-def build_network(net_name, ae_net=None, rep_dim=1024):
+def build_network(net_name, rep_dim, ae_net=None):
     """Builds the neural network."""
 
     implemented_networks = (
@@ -129,7 +129,7 @@ def build_network(net_name, ae_net=None, rep_dim=1024):
     return net
 
 
-def build_autoencoder(net_name):
+def build_autoencoder(net_name, rep_dim):
     """Builds the corresponding autoencoder network."""
 
     implemented_networks = (
@@ -158,7 +158,7 @@ def build_autoencoder(net_name):
         ae_net = MNIST_LeNet_Autoencoder()
 
     if net_name == "subter_LeNet":
-        ae_net = SubTer_LeNet_Autoencoder()
+        ae_net = SubTer_LeNet_Autoencoder(rep_dim=rep_dim)
 
     if net_name == "subter_LeNet_Split":
         ae_net = SubTer_LeNet_Split_Autoencoder()

Deep-SAD-PyTorch/src/optim/ae_trainer.py

@@ -85,6 +85,7 @@ class AETrainer(BaseTrainer):
         logger.info("Starting pretraining...")
         start_time = time.time()
         ae_net.train()
+        ae_net.summary(receptive_field=True)  # Add network summary before training
 
         for epoch in range(self.n_epochs):
             epoch_loss = 0.0
@@ -197,6 +198,8 @@ class AETrainer(BaseTrainer):
         n_batches = 0
         start_time = time.time()
         ae_net.eval()
+        ae_net.summary(receptive_field=True)  # Add network summary before testing
+
         with torch.no_grad():
             for data in test_loader:
                 (