Merge pull request #98 from ispras/extend_models_stage3

Extend models stage3

Author: LukyanovKirillML

experiments/attack_defense_exps.py

@@ -3,12 +3,13 @@
 import torch
 from torch import device
 
-from data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
-from datasets.datasets_manager import DatasetManager
-from datasets.ptg_datasets import LibPTGDataset
-from models_builder.gnn_models import FrameworkGNNModelManager, Metric
-from models_builder.models_zoo import model_configs_zoo
-from aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, \
+from gnn_aid.data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
+from gnn_aid.datasets.datasets_manager import DatasetManager
+from gnn_aid.datasets.ptg_datasets import LibPTGDataset
+from gnn_aid.models_builder import Metric
+from gnn_aid.models_builder.model_managers import FrameworkGNNModelManager
+from gnn_aid.models_builder.models_zoo import model_configs_zoo
+from gnn_aid.aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, \
     EVASION_ATTACK_PARAMETERS_PATH, \
     EVASION_DEFENSE_PARAMETERS_PATH
 
@@ -758,8 +759,8 @@ def test_adv_training():
             # "num_nodes": dataset.dataset.x.shape[0]
         }
     )
-    from defenses.evasion_defense import EvasionDefender
-    from aux.utils import all_subclasses
+    from gnn_aid.defenses.evasion_defense import EvasionDefender
+    from gnn_aid.aux.utils import all_subclasses
     print([e.name for e in all_subclasses(EvasionDefender)])
     gnn_model_manager.set_evasion_defender(evasion_defense_config=evasion_defense_config)
 
@@ -1375,6 +1376,7 @@ def test_rewatt():
         DatasetConfig(full_name),
         LibPTGDataset.default_dataset_var_config.clone_with({"task": Task.NODE_CLASSIFICATION})
     )
+    data = dataset.data
     data.to(my_device)
 
     gcn_gcn = model_configs_zoo(dataset=dataset, model_name='gcn_gcn')

experiments/attack_defense_metric_test.py

@@ -4,15 +4,15 @@
 import torch
 from torch import device
 
-from datasets.datasets_manager import DatasetManager
-from datasets.ptg_datasets import LibPTGDataset
-from models_builder.attack_defense_manager import FrameworkAttackDefenseManager
-from models_builder.attack_defense_metric import AttackMetric, DefenseMetric
-from aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, EVASION_ATTACK_PARAMETERS_PATH, \
+from gnn_aid.datasets.datasets_manager import DatasetManager
+from gnn_aid.datasets.ptg_datasets import LibPTGDataset
+from gnn_aid.models_builder.attack_defense_manager import FrameworkAttackDefenseManager
+from gnn_aid.models_builder.attack_defense_metric import AttackMetric, DefenseMetric
+from gnn_aid.aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, EVASION_ATTACK_PARAMETERS_PATH, \
     EVASION_DEFENSE_PARAMETERS_PATH
-from models_builder.gnn_models import FrameworkGNNModelManager
-from data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
-from models_builder.models_zoo import model_configs_zoo
+from gnn_aid.models_builder.model_managers import FrameworkGNNModelManager
+from gnn_aid.data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
+from gnn_aid.models_builder.models_zoo import model_configs_zoo
 
 
 def attack_defense_metrics():

experiments/attack_defense_test.py

@@ -4,20 +4,21 @@
 
 from torch import device
 
-from datasets.ptg_datasets import LibPTGDataset
-from models_builder.models_utils import apply_decorator_to_graph_layers
-from aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, EVASION_ATTACK_PARAMETERS_PATH, \
+from gnn_aid.datasets.ptg_datasets import LibPTGDataset
+from gnn_aid.models_builder.models_utils import apply_decorator_to_graph_layers
+from gnn_aid.aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, EVASION_ATTACK_PARAMETERS_PATH, \
     EVASION_DEFENSE_PARAMETERS_PATH
-from models_builder.gnn_models import FrameworkGNNModelManager, Metric
-from data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
-from datasets.datasets_manager import DatasetManager
-from models_builder.models_zoo import model_configs_zoo
-from attacks.qattack import qattack
+from gnn_aid.models_builder.model_managers import FrameworkGNNModelManager
+from gnn_aid.models_builder.models_utils import Metric
+from gnn_aid.data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
+from gnn_aid.datasets.datasets_manager import DatasetManager
+from gnn_aid.models_builder.models_zoo import model_configs_zoo
+from gnn_aid.attacks.qattack import qattack
 # from attacks.RL_S2V.rl_s2v import RLS2VAttacker
-from defenses.jaccard_defense import jaccard_def
-from attacks.metattack import meta_gradient_attack
-from defenses.gnn_guard import gnnguard
-from defenses.pro_gnn.prognn import ProGNNDefender
+from gnn_aid.defenses.jaccard_defense import jaccard_def
+from gnn_aid.attacks.metattack import meta_gradient_attack
+from gnn_aid.defenses.gnn_guard import gnnguard
+from gnn_aid.defenses.pro_gnn.prognn import ProGNNDefender
 
 
 def test_attack_defense_small():
@@ -206,7 +207,7 @@ def test_attack_defense_small():
 
 
 def test_attack_defense():
-    from attacks.clga import CLGA
+    from gnn_aid.attacks.clga import CLGA
 
     my_device = device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -980,8 +981,8 @@ def test_adv_training():
             # "num_nodes": dataset.dataset.x.shape[0]
         }
     )
-    from defenses.evasion_defense import EvasionDefender
-    from aux.utils import all_subclasses
+    from gnn_aid.defenses.evasion_defense import EvasionDefender
+    from gnn_aid.aux.utils import all_subclasses
     print([e.name for e in all_subclasses(EvasionDefender)])
     gnn_model_manager.set_evasion_defender(evasion_defense_config=evasion_defense_config)
 

experiments/backend_demo_ase.py

@@ -5,12 +5,13 @@
 
 from torch import device
 
-from aux.utils import EVASION_ATTACK_PARAMETERS_PATH, EVASION_DEFENSE_PARAMETERS_PATH
-from datasets.ptg_datasets import LibPTGDataset
-from models_builder.gnn_models import FrameworkGNNModelManager, Metric
-from data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
-from datasets.datasets_manager import DatasetManager
-from models_builder.models_zoo import model_configs_zoo
+from gnn_aid.aux.utils import EVASION_ATTACK_PARAMETERS_PATH, EVASION_DEFENSE_PARAMETERS_PATH
+from gnn_aid.datasets.ptg_datasets import LibPTGDataset
+from gnn_aid.models_builder.model_managers import FrameworkGNNModelManager
+from gnn_aid.models_builder.models_utils import Metric
+from gnn_aid.data_structures.configs import ModelModificationConfig, ConfigPattern, DatasetConfig, Task
+from gnn_aid.datasets.datasets_manager import DatasetManager
+from gnn_aid.models_builder.models_zoo import model_configs_zoo
 
 
 def test_attack_defense_small():

experiments/user_dataset.py

@@ -2,7 +2,7 @@
 from torch import tensor
 from torch_geometric.data import InMemoryDataset, Data, Dataset
 
-from datasets.datasets_manager import DatasetManager
+from gnn_aid.datasets.datasets_manager import DatasetManager
 
 
 # Example of local user PTG dataset

experiments/various_tasks.py

@@ -1,13 +1,18 @@
+import json
+
 import torch
 from torch import device
 
 from gnn_aid.attacks import Attacker
 from gnn_aid.aux.utils import FUNCTIONS_PARAMETERS_PATH, all_subclasses
+from gnn_aid.data_structures import ModelStructureConfig, ModelConfig
 from gnn_aid.data_structures.configs import DatasetConfig, DatasetVarConfig, FeatureConfig, Task, \
     ConfigPattern, ModelModificationConfig
 from gnn_aid.datasets.datasets_manager import DatasetManager
 from gnn_aid.datasets.ptg_datasets import LibPTGDataset
-from gnn_aid.models_builder.gnn_models import FrameworkGNNModelManager, Metric
+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.models_builder.models_zoo import model_configs_zoo
 
 
@@ -103,9 +108,32 @@ def link_prediction():
 
     gen_dataset = DatasetManager.get_by_config(dc, dvc)
     print(gen_dataset.data)
-    gen_dataset.train_test_split()
+    gen_dataset.train_test_split(percent_train_class=0.85, percent_test_class=0.1)
+
+    gnn = FrameworkGNNConstructor(
+        model_config=ModelConfig(
+            structure=ModelStructureConfig(
+                [
+                    {
+                        'label': 'n',
+                        'layer': {
+                            'layer_name': 'SAGEConv',
+                            'layer_kwargs': {
+                                'in_channels': gen_dataset.num_node_features,
+                                'out_channels': 16,
+                            },
+                        },
+                    },
+                    {
+                        'label': 'd',
+                        'function': {
+                            'function_name': 'CosineSimilarity',
+                            'function_kwargs': None
+                        }
+                    }
+                ]
+            )))
 
-    gnn = model_configs_zoo(dataset=gen_dataset, model_name='gcn_gcn')
     manager_config = ConfigPattern(
         _config_class="ModelManagerConfig",
         _config_kwargs={
@@ -125,7 +153,7 @@ def link_prediction():
         }
     )
 
-    steps_epochs = 30
+    steps_epochs = 3
     my_device = device('cuda' if torch.cuda.is_available() else 'cpu')
     gnn_model_manager = FrameworkGNNModelManager(
         gnn=gnn,
@@ -145,13 +173,130 @@ def link_prediction():
     )
     print("Training was successful")
 
+    res = gnn_model_manager.run_model(
+        gen_dataset=gen_dataset,
+        mask='all'
+    )
+    res = gnn_model_manager.evaluate_model(
+        gen_dataset=gen_dataset,
+        metrics=[Metric("Accuracy", mask="test")]
+    )
+    print(json.dumps(res, indent=2))
+
+    # gnn.eval()
+    # z = gnn.encode(data.x, data.edge_index)
+    # out = gnn.decode(z, data.edge_label_index).view(-1).sigmoid()
+    # return roc_auc_score(data.edge_label.cpu().numpy(), out.cpu().numpy())
+
+
+def ptg_example():
+    import os.path as osp
+
+    import torch
+    from sklearn.metrics import roc_auc_score
+
+    import torch_geometric.transforms as T
+    from torch_geometric.datasets import Planetoid
+    from torch_geometric.nn import GCNConv
+    from torch_geometric.utils import negative_sampling
+
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        device = torch.device('mps')
+    else:
+        device = torch.device('cpu')
+
+    transform = T.Compose([
+        T.NormalizeFeatures(),
+        T.ToDevice(device),
+        T.RandomLinkSplit(num_val=0.05, num_test=0.1, is_undirected=True,
+                          add_negative_train_samples=False),
+    ])
+    path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Planetoid')
+    dataset = Planetoid(path, name='Cora', transform=transform)
+    # After applying the `RandomLinkSplit` transform, the data is transformed from
+    # a data object to a list of tuples (train_data, val_data, test_data), with
+    # each element representing the corresponding split.
+    train_data, val_data, test_data = dataset[0]
+
+    class Net(torch.nn.Module):
+        def __init__(self, in_channels, hidden_channels, out_channels):
+            super().__init__()
+            self.conv1 = GCNConv(in_channels, hidden_channels)
+            self.conv2 = GCNConv(hidden_channels, out_channels)
+
+        def encode(self, x, edge_index):
+            x = self.conv1(x, edge_index).relu()
+            return self.conv2(x, edge_index)
+
+        def decode(self, z, edge_label_index):
+            return (z[edge_label_index[0]] * z[edge_label_index[1]]).sum(dim=-1)
+
+        def decode_all(self, z):
+            prob_adj = z @ z.t()
+            return (prob_adj > 0).nonzero(as_tuple=False).t()
+
+    model = Net(dataset.num_features, 128, 64).to(device)
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+    def train():
+        model.train()
+        optimizer.zero_grad()
+        z = model.encode(train_data.x, train_data.edge_index)
+
+        # We perform a new round of negative sampling for every training epoch:
+        neg_edge_index = negative_sampling(
+            edge_index=train_data.edge_index, num_nodes=train_data.num_nodes,
+            num_neg_samples=train_data.edge_label_index.size(1), method='sparse')
+
+        edge_label_index = torch.cat(
+            [train_data.edge_label_index, neg_edge_index],
+            dim=-1,
+        )
+        edge_label = torch.cat([
+            train_data.edge_label,
+            train_data.edge_label.new_zeros(neg_edge_index.size(1))
+        ], dim=0)
+
+        out = model.decode(z, edge_label_index).view(-1)
+        loss = criterion(out, edge_label)
+        loss.backward()
+        optimizer.step()
+        return loss
+
+    @torch.no_grad()
+    def test(data):
+        model.eval()
+        z = model.encode(data.x, data.edge_index)
+        out = model.decode(z, data.edge_label_index).view(-1).sigmoid()
+        return roc_auc_score(data.edge_label.cpu().numpy(), out.cpu().numpy())
+
+    best_val_auc = final_test_auc = 0
+    for epoch in range(1, 101):
+        loss = train()
+        val_auc = test(val_data)
+        test_auc = test(test_data)
+        if val_auc > best_val_auc:
+            best_val_auc = val_auc
+            final_test_auc = test_auc
+        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Val: {val_auc:.4f}, '
+              f'Test: {test_auc:.4f}')
+
+    print(f'Final Test: {final_test_auc:.4f}')
+
+    z = model.encode(test_data.x, test_data.edge_index)
+    final_edge_index = model.decode_all(z)
+
 
 if __name__ == '__main__':
     # node_regression()
     # graph_regression()
 
     # edge_regression()
-    # link_prediction()
+    link_prediction()
 
-    for c in all_subclasses(Attacker):
-        print(c)
+    # ptg_example()
+    # for c in all_subclasses(Attacker):
+    #     print(c)

gnn_aid/attacks/attack_base.py

@@ -1,6 +1,6 @@
 from gnn_aid.datasets.gen_dataset import GeneralDataset
 from gnn_aid.data_structures.graph_modification_artifacts import GraphModificationArtifact
-from gnn_aid.models_builder.gnn_models import GNNModelManager
+from gnn_aid.models_builder.model_managers import GNNModelManager
 
 
 class Attacker:

gnn_aid/attacks/evasion_attacks.py

@@ -5,7 +5,7 @@
 from gnn_aid.attacks.attack_base import Attacker
 from gnn_aid.aux.utils import move_to_same_device
 from gnn_aid.datasets.gen_dataset import GeneralDataset
-from gnn_aid.models_builder.gnn_models import GNNModelManager
+from gnn_aid.models_builder.model_managers import GNNModelManager
 
 # Nettack imports
 from .evasion_attacks_collection.nettack.utils import NettackSurrogate, NettackAttack

gnn_aid/attacks/metattack/meta_gradient_attack.py

@@ -13,7 +13,7 @@
 from gnn_aid.aux.utils import OPTIMIZERS_PARAMETERS_PATH, move_to_same_device
 from gnn_aid.data_structures.configs import ModelModificationConfig, ConfigPattern
 from gnn_aid.datasets.gen_dataset import GeneralDataset
-from gnn_aid.models_builder.gnn_models import FrameworkGNNModelManager, GNNModelManager
+from gnn_aid.models_builder.model_managers import GNNModelManager, FrameworkGNNModelManager
 from gnn_aid.models_builder.models_zoo import model_configs_zoo
 from . import utils
 

gnn_aid/data_structures/configs.py

@@ -616,7 +616,7 @@ def __init__(
     @property
     def task(
             self
-    ) -> Union[str, dict]:
+    ) -> Task:
         return self["task"]
 
     @property