Add CLGA LP test

Author: Jeratt

tests/attacks_test.py

@@ -1,17 +1,20 @@
+import copy
 import unittest
 import numpy as np
 import torch
 
+from gnn_aid.attacks.clga.CLGA import CLGAAttack
 from gnn_aid.attacks.mi_attacks import MIAttacker
 from gnn_aid.datasets.datasets_manager import DatasetManager
 from gnn_aid.datasets.ptg_datasets import LibPTGDataset
+from gnn_aid.models_builder import FrameworkGNNConstructor
 from gnn_aid.models_builder.models_utils import Metric
 from gnn_aid.models_builder.model_managers import FrameworkGNNModelManager
 from gnn_aid.data_structures.configs import ModelModificationConfig, DatasetConfig, DatasetVarConfig, \
-    ConfigPattern, FeatureConfig, Task
+    ConfigPattern, FeatureConfig, Task, ModelConfig, ModelStructureConfig
 from gnn_aid.models_builder.models_zoo import model_configs_zoo
 from gnn_aid.aux.utils import POISON_ATTACK_PARAMETERS_PATH, EVASION_ATTACK_PARAMETERS_PATH, \
-    OPTIMIZERS_PARAMETERS_PATH, MI_ATTACK_PARAMETERS_PATH
+    OPTIMIZERS_PARAMETERS_PATH, MI_ATTACK_PARAMETERS_PATH, FUNCTIONS_PARAMETERS_PATH
 from .utils import monkey_patch_dirs, cleanup_patches
 
 
@@ -695,6 +698,118 @@ def test_mi_shadow_cora(self):
             print(f"MI Attack accuracy:"
                   f" {MIAttacker.compute_single_attack_accuracy(mask, res, self.gen_dataset_sg_cora.train_mask)}")
 
+    def test_clga_link_prediction(self):
+        gen_dataset = DatasetManager.get_by_config(
+            DatasetConfig((LibPTGDataset.data_folder, "Homogeneous", "Planetoid", "Cora")),
+            LibPTGDataset.default_dataset_var_config.clone_with({"task": Task.EDGE_PREDICTION})
+        )
+
+        poison_attack_config = ConfigPattern(
+            _class_name="CLGAAttack",
+            _import_path=POISON_ATTACK_PARAMETERS_PATH,
+            _config_class="PoisonAttackConfig",
+            _config_kwargs={
+                "learning_rate": 0.01,
+                "num_epochs": 50,
+            }
+        )
+
+        gnn = FrameworkGNNConstructor(
+            model_config=ModelConfig(
+                structure=ModelStructureConfig([
+                    {
+                        'label': 'n',
+                        'layer': {
+                            'layer_name': 'GCNConv',
+                            'layer_kwargs': {'in_channels': gen_dataset.num_node_features, 'out_channels': 32}
+                        },
+                        'activation': {
+                            'activation_name': 'ReLU',
+                            'activation_kwargs': None
+                        }
+                    },
+                    {
+                        'label': 'n',
+                        'layer': {
+                            'layer_name': 'GCNConv',
+                            'layer_kwargs': {'in_channels': 32, 'out_channels': 16}
+                        }
+                    },
+                    {
+                        'label': 'd',
+                        'function': {
+                            'function_name': 'Concat',
+                            'function_kwargs': None
+                        }
+                    },
+                    {
+                        'label': 'd',
+                        'layer': {
+                            'layer_name': 'Linear',
+                            'layer_kwargs': {'in_features': 32, 'out_features': 16}
+                        },
+                        'activation': {
+                            'activation_name': 'ReLU',
+                            'activation_kwargs': None
+                        }
+                    },
+                    {
+                        'label': 'd',
+                        'layer': {
+                            'layer_name': 'Linear',
+                            'layer_kwargs': {'in_features': 16, 'out_features': 1}
+                        }
+                    }
+                ])
+            )
+        )
+
+        manager_config_lp = ConfigPattern(
+            _config_class="ModelManagerConfig",
+            _config_kwargs={
+                "mask_features": [],
+                "optimizer": {
+                    "_class_name": "Adam",
+                    "_config_kwargs": {
+                        "lr": 0.01,
+                        "weight_decay": 5e-4
+                    },
+                },
+                "loss_function": {
+                    "_class_name": "BCEWithLogitsLoss",
+                    "_import_path": FUNCTIONS_PARAMETERS_PATH,
+                    "_class_import_info": ["torch.nn"],
+                    "_config_kwargs": {}
+                },
+                "batch": 64
+            }
+        )
+
+        gnn_model_manager = FrameworkGNNModelManager(
+            gnn=gnn,
+            dataset_path=gen_dataset.prepared_dir,
+            modification=ModelModificationConfig(model_ver_ind=0, epochs=30),
+            manager_config=manager_config_lp,
+        )
+
+        gnn_model_manager.set_poison_attacker(poison_attack_config=poison_attack_config)
+
+        gen_dataset.train_test_split(percent_train_class=0.85, percent_test_class=0.15)
+
+        gnn_model_manager.train_model(
+            gen_dataset=gen_dataset,
+            steps=30,
+            metrics=[Metric("AUC", mask='train')]
+        )
+
+        test_metrics = gnn_model_manager.evaluate_model(
+            gen_dataset=gen_dataset,
+            metrics=[Metric("AUC", mask='test'), Metric("Recall@k", mask='test', k=100)]
+        )
+        print("CLGA Link Prediction AUC:", test_metrics['test']['AUC'])
+
+        self.assertLess(test_metrics['test']['AUC'], 0.95)
+
 
 if __name__ == '__main__':
     unittest.main()