ready

eval/UnlearningAttack.py

@@ -164,7 +164,8 @@ class UnlearningAttack:
             naive_retrained, 
             test_loader, 
             device, 
-            target_class
+            target_class, 
+            strategy_in_use,
         ):
 
         filtered_model.eval()
@@ -184,6 +185,13 @@ class UnlearningAttack:
                     out_filtered = filtered_model(data)
                 
                 out_naive = naive_retrained(data)
+
+                if(strategy_in_use == "LinearFiltration"):
+                    mask = torch.ones(out_filtered.shape[1], dtype = torch.bool, device = device)
+                    mask[target_class] = False
+                    # out
+                    out_filtered = out_filtered[:,mask]
+                    out_naive = out_naive[:, mask]
                 
                 filtered_logits.append(out_filtered)
                 naive_logits.append(out_naive)
@@ -206,7 +214,8 @@ class UnlearningAttack:
             test_loader, 
             unlearned_instance, 
             base_shadow_instance, 
-            device
+            device,
+            strategy_in_use
         ):
         
 
@@ -217,7 +226,7 @@ class UnlearningAttack:
         
         if not os.path.exists(current_log_file):
             with open(current_log_file, "w") as f:
-                f.write("target_class, parameter_mia_accuracy, latent_distance_tell, lookalike_accuracy, A-Dist, JS-Dist\n")
+                f.write("target_class, parameter_mia_accuracy, lookalike_accuracy, A-Dist, JS-Dist\n")
 
         # Parameter-Space MIA 
         parameter_mia_acc =  self.run_parameter_space_mia(
@@ -248,7 +257,8 @@ class UnlearningAttack:
             naive_retrained=reference_model_torch, 
             test_loader=test_loader, 
             device=device,
-            target_class=target_class
+            target_class=target_class,
+            strategy_in_use = strategy_in_use
         )
 
         # Calculate JS-Dist