Bug fixes

Author: projectavi

classification_pipeline.py

@@ -41,7 +41,7 @@ def construct_classifier(X: torch.tensor, t: torch.tensor, k: int, model_config:
     training_config.set_logits_loss(logits_loss)
 
     # Create a temporary target vector where everything equal to i is 1 and the rest are 0
-    temp_target = torch.tensor([1 if x == k else 0 for x in t])
+    temp_target = t
 
     # Training loop
     model, metrics = train_model(model, X, temp_target, training_config)
@@ -141,8 +141,12 @@ def train_model(model: nn.Module, X: torch.tensor, t: torch.tensor, training_con
             train_accuracies.append(accuracy(outputs, labels, logits_loss))
             val_accuracies.append(accuracy(model(X_val), t_val, logits_loss))
 
-        print("Epoch ", epoch, " Train Loss: ", train_losses[-1], " Train Accuracy: ", train_accuracies[-1],
-              " Val Loss: ", val_losses[-1], " Val Accuracy: ", val_accuracies[-1])
+        full_train_loss = criterion(model(X_train), t_train).item()
+        full_val_loss = criterion(model(X_val), t_val).item()
+        full_train_accuracy = accuracy(model(X_train), t_train, logits_loss)
+        full_val_accuracy = accuracy(model(X_val), t_val, logits_loss)
+        print("Epoch ", epoch, " Train Loss: ", full_train_loss, " Train Accuracy: ", full_train_accuracy,
+              " Val Loss: ", full_val_loss, " Val Accuracy: ", full_val_accuracy)
 
     metrics = {
         "train_losses": train_losses,

main.py

@@ -32,58 +32,60 @@
 
     # Construct the set of unique labels
     targets = [x.item() for x in labels]
-    unique_labels = set(targets)
+    unique_labels = list(set(targets))
+
+    model_format = {
+        "num_layers": 4,
+        "layers": [
+            {
+                "layer_type": "Linear_LRP",
+                "in_dim": 50,
+                "out_dim": 64,
+                "act": "ReLU",
+                "dropout": 0.1,
+                "batch_norm": False
+            },
+            {
+                "layer_type": "Linear_LRP",
+                "in_dim": 64,
+                "out_dim": 128,
+                "act": "ReLU",
+                "dropout": 0.0,
+                "batch_norm": False
+            },
+            {
+                "layer_type": "Linear_LRP",
+                "in_dim": 128,
+                "out_dim": 32,
+                "act": "ReLU",
+                "dropout": 0.0,
+                "batch_norm": False
+            },
+            {
+                "layer_type": "Linear_LRP",
+                "in_dim": 32,
+                "out_dim": 1,
+                "act": "Sigmoid",
+                "dropout": 0.0,
+                "batch_norm": False
+            },
+        ]
+    }
+
+    model_config = ModelConfig(model_format)
 
     # Check if the models have been saved
     if os.path.exists("models.pkl"):
         models = pickle.load(open("models.pkl", "rb"))
         print("Models loaded")
     else:
         print("Models not found")
-        model_format = {
-            "num_layers": 4,
-            "layers": [
-                {
-                    "layer_type": "Linear_LRP",
-                    "in_dim": 50,
-                    "out_dim": 64,
-                    "act": "ReLU",
-                    "dropout": 0.1,
-                    "batch_norm": False
-                },
-                {
-                    "layer_type": "Linear_LRP",
-                    "in_dim": 64,
-                    "out_dim": 128,
-                    "act": "ReLU",
-                    "dropout": 0.0,
-                    "batch_norm": False
-                },
-                {
-                    "layer_type": "Linear_LRP",
-                    "in_dim": 128,
-                    "out_dim": 32,
-                    "act": "ReLU",
-                    "dropout": 0.0,
-                    "batch_norm": False
-                },
-                {
-                    "layer_type": "Linear_LRP",
-                    "in_dim": 32,
-                    "out_dim": 1,
-                    "act": "Sigmoid",
-                    "dropout": 0.0,
-                    "batch_norm": False
-                },
-            ]
-        }
-
-        model_config = ModelConfig(model_format)
 
         training_config = TrainingConfig(0.01, 1, batch_size)
 
         models = {}
 
+        # k = unique_labels[0]
         for k in unique_labels:
             print("Training model for class", k)
             # Create a copy of X
@@ -107,23 +109,29 @@
         explanations = pickle.load(open("explanations.pkl", "rb"))
         print("Explanations loaded")
     else:
+        X.requires_grad = True
+
         explanations = {}
 
+        if model_config.layers[-1]["act"] == "Sigmoid":
+            criterion = nn.BCELoss()
+        else:
+            criterion = nn.BCEWithLogitsLoss()
+
+        # k = unique_labels[0]
         for k in unique_labels:
             print("Explaining model for class", k)
             model = models[k]
 
             predictions = model.forward(X, explain=True, rule="alpha2beta1")
 
-            predictions = predictions[torch.arange(batch_size), predictions.max(1)[1]]  # Choose maximizing output neuron
-
-            predictions = predictions.sum()
+            pred = predictions.sum()
 
-            predictions.backward()
+            loss = criterion(predictions, torch.reshape(torch.tensor([1 if x == k else 0 for x in labels]), (-1, 1)).float())
 
-            explanation = X.grad
+            pred.backward()
 
-            explanations[k] = explanation
+            explanations[k] = X.grad
 
             # Save the explanations
             pickle.dump(explanations, open("explanations.pkl", "wb"))

model_config.py

@@ -116,12 +116,15 @@ def __init__(self, modelConfig, lrp_bool=False):
 
         self.layers = nn.ModuleList(self.layers)
 
-    def forward(self, x):
+    def forward(self, x, explain=False, rule="epsilon", pattern=None):
         # assume X is already in the correct shape and all that
         out = x
 
         for layer in self.layers:
-            out = layer(out)
+            if explain:
+                out = layer.forward(out, explain=explain, rule=rule, pattern=pattern)
+            else:
+                out = layer(out)
 
         return out