prepare dataset for another tasks

Author: mishadr

data/example/example/metainfo

@@ -18,6 +18,7 @@
  },
  "labelings": {
   "binary": 2,
-  "threeClasses": 3
+  "threeClasses": 3,
+  "regression": 0
  }
 }

data/example/example/raw/labels/regression

@@ -0,0 +1,11 @@
+{
+"10": 0.1,
+"11": 0.2,
+"12": 1.1,
+"13": 1.2,
+"14": 2.5,
+"15": 2.3,
+"16": 2.1,
+"17": 2.0
+}
+

experiments/various_tasks.py

@@ -0,0 +1,54 @@
+import torch
+from torch import device
+
+from data_structures.configs import DatasetConfig, DatasetVarConfig, FeatureConfig, Task, \
+    ConfigPattern, ModelModificationConfig
+from datasets.datasets_manager import DatasetManager
+from models_builder.gnn_models import FrameworkGNNModelManager, Metric
+from models_builder.models_zoo import model_configs_zoo
+
+
+def regression():
+    dc = DatasetConfig(('example', 'example'))
+    dvc = DatasetVarConfig(features=FeatureConfig(node_attr=['a']), labeling='regression',
+                           task=Task.NODE_REGRESSION, dataset_ver_ind=0)
+
+    gen_dataset = DatasetManager.get_by_config(dc, dvc)
+
+    print(gen_dataset.data)
+
+    gnn = model_configs_zoo(dataset=gen_dataset, model_name='gcn_gcn')
+    manager_config = ConfigPattern(
+        _config_class="ModelManagerConfig",
+        _config_kwargs={
+            "mask_features": [],
+            "optimizer": {
+                "_class_name": "Adam",
+                "_config_kwargs": {},
+            }
+        }
+    )
+
+    steps_epochs = 10
+    my_device = device('cuda' if torch.cuda.is_available() else 'cpu')
+    gnn_model_manager = FrameworkGNNModelManager(
+        gnn=gnn,
+        dataset_path=gen_dataset.prepared_dir,
+        manager_config=manager_config,
+        modification=ModelModificationConfig(model_ver_ind=0, epochs=steps_epochs)
+    )
+
+    gnn_model_manager.gnn.to(my_device)
+    gen_dataset.data.to(my_device)
+
+    gen_dataset.train_test_split()
+    gnn_model_manager.train_model(
+        gen_dataset=gen_dataset, steps=steps_epochs,
+        save_model_flag=False,
+        metrics=[Metric("F1", mask='train', average=None)]
+    )
+    print("Training was successful")
+
+
+if __name__ == '__main__':
+    regression()

src/data_structures/configs.py

@@ -2,6 +2,7 @@
 import inspect
 import json
 import logging
+from enum import Enum
 from json import JSONEncoder
 from pathlib import Path
 from typing import Union, Any, Type, Tuple, Self
@@ -17,6 +18,13 @@
 DATA_CHANGE_FLAG = "__data_change_flag"
 
 
+class Task(str, Enum):
+    NODE_CLASSIFICATION = "NODE_CLASSIFICATION"
+    GRAPH_CLASSIFICATION = "GRAPH_CLASSIFICATION"
+    NODE_REGRESSION = "NODE_REGRESSION"
+    LINK_PREDICTION = "LINK_PREDICTION"
+
+
 # TECHNICAL_KEYS_SET_FOR_CONFIGS = {CONFIG_PARAMS_PATH_KEY, CONFIG_CLASS_NAME,
 #                                   CONFIG_SAVE_KWARGS_KEY, DATA_CHANGE_FLAG}
 
@@ -572,13 +580,13 @@ def __init__(
             self,
             features: FeatureConfig = None,
             labeling: Union[str, dict] = None,
-            # task: str = None,
+            task: Task = None,
             dataset_ver_ind: int = None,
             **kwargs
     ):
         """ """
         super().__init__(
-            features=features, labeling=labeling, dataset_ver_ind=dataset_ver_ind, **kwargs)
+            features=features, labeling=labeling, task=task, dataset_ver_ind=dataset_ver_ind, **kwargs)
 
     @property
     def features(
@@ -592,6 +600,12 @@ def labeling(
     ) -> Union[str, dict]:
         return self["labeling"]
 
+    @property
+    def task(
+            self
+    ) -> Union[str, dict]:
+        return self["task"]
+
     @property
     def dataset_ver_ind(
             self
@@ -873,7 +887,7 @@ class ModelModificationConfig(
 
     def __init__(
             self,
-            model_ver_ind: [int, None] = None,
+            model_ver_ind: int | None = None,
             epochs=None,
             **kwargs
     ):

src/datasets/dataset_info.py

@@ -143,7 +143,7 @@ def check_validity(
             labelings = list(self.labelings.items())
         for k, v in labelings:
             assert isinstance(k, str)
-            assert isinstance(v, int) and v >= 1  # 1 stands for regression
+            assert isinstance(v, int) and v >= 0  # 1 stands for regression
 
     def check_consistency(
             self

src/datasets/gen_dataset.py

@@ -10,7 +10,8 @@
 
 from aux.declaration import Declare
 from aux.utils import root_dir
-from data_structures.configs import DatasetConfig, DatasetVarConfig, ConfigPattern, FeatureConfig
+from data_structures.configs import DatasetConfig, DatasetVarConfig, ConfigPattern, FeatureConfig, \
+    Task
 from datasets.dataset_info import DatasetInfo
 from datasets.visible_part import VisiblePart
 
@@ -297,25 +298,33 @@ def train_test_split(
 
         if percent_val_class < -1.1e-15:
             raise RuntimeError("percent_train_class + percent_test_class > 1")
-        train_mask = torch.BoolTensor([False] * self.labels.size(dim=0))
-        val_mask = torch.BoolTensor([False] * self.labels.size(dim=0))
-        test_mask = torch.BoolTensor([False] * self.labels.size(dim=0))
-
-        labeled_nodes_numbers = [n for n, y in enumerate(self.labels) if y != -1]
-        num_train = int(percent_train_class * len(labeled_nodes_numbers))
-        num_test = int(percent_test_class * len(labeled_nodes_numbers))
-        num_eval = len(labeled_nodes_numbers) - num_train - num_test
-        if percent_val_class <= 0 and num_eval > 0:
-            num_test += num_eval
-            num_eval = 0
-        split = randperm(num_train + num_eval + num_test, generator=default_generator).tolist()
-
-        for elem in split[:num_train]:
-            train_mask[labeled_nodes_numbers[elem]] = True
-        for elem in split[num_train: num_train + num_eval]:
-            val_mask[labeled_nodes_numbers[elem]] = True
-        for elem in split[num_train + num_eval:]:
-            test_mask[labeled_nodes_numbers[elem]] = True
+
+        task_type = self.dataset_var_config.task
+        if task_type in [Task.NODE_CLASSIFICATION, Task.NODE_REGRESSION, Task.GRAPH_CLASSIFICATION]:
+            train_mask = torch.BoolTensor([False] * self.labels.size(dim=0))
+            val_mask = torch.BoolTensor([False] * self.labels.size(dim=0))
+            test_mask = torch.BoolTensor([False] * self.labels.size(dim=0))
+
+            labeled_nodes_numbers = [n for n, y in enumerate(self.labels) if y != -1]
+            num_train = int(percent_train_class * len(labeled_nodes_numbers))
+            num_test = int(percent_test_class * len(labeled_nodes_numbers))
+            num_eval = len(labeled_nodes_numbers) - num_train - num_test
+            if percent_val_class <= 0 and num_eval > 0:
+                num_test += num_eval
+                num_eval = 0
+            split = randperm(num_train + num_eval + num_test, generator=default_generator).tolist()
+
+            for elem in split[:num_train]:
+                train_mask[labeled_nodes_numbers[elem]] = True
+            for elem in split[num_train: num_train + num_eval]:
+                val_mask[labeled_nodes_numbers[elem]] = True
+            for elem in split[num_train + num_eval:]:
+                test_mask[labeled_nodes_numbers[elem]] = True
+
+        elif task_type == Task.LINK_PREDICTION:
+            raise NotImplementedError
+        else:
+            raise ValueError(f"Unsupported task type {task_type}")
 
         self.train_mask = train_mask
         self.test_mask = test_mask

src/models_builder/gnn_models.py

@@ -24,7 +24,7 @@
 from data_structures.configs import ConfigPattern, PoisonAttackConfig, CONFIG_OBJ, \
     EvasionAttackConfig, MIAttackConfig, PoisonDefenseConfig, EvasionDefenseConfig, \
     MIDefenseConfig, ModelManagerConfig, ModelModificationConfig, ModelConfig, \
-    CONFIG_CLASS_NAME
+    CONFIG_CLASS_NAME, Task
 from data_structures.graph_modification_artifacts import GraphModificationArtifact
 from datasets.gen_dataset import GeneralDataset
 from web_interface.back_front.utils import SocketConnect
@@ -946,9 +946,8 @@ def train_1_step(
             self,
             gen_dataset: GeneralDataset
     ) -> List[Union[float, int]]:
-        # FIXME misha it is not task type, change to getting dvc field
-        task_type = "multiple-graphs" if gen_dataset.is_multi() else "single-graph"
-        if task_type == "single-graph":
+        task_type = gen_dataset.dataset_var_config.task
+        if task_type == Task.NODE_CLASSIFICATION:
             # FIXME Kirill, add data_x_copy mask
             loader = cast(
                 Iterable,
@@ -958,7 +957,7 @@ def train_1_step(
                     batch_size=self.batch, shuffle=True
                 )
             )
-        elif task_type == "multiple-graphs":
+        elif task_type == Task.GRAPH_CLASSIFICATION:
             train_dataset = gen_dataset.dataset.index_select(gen_dataset.train_mask)
             loader = cast(
                 Iterable,
@@ -967,7 +966,7 @@ def train_1_step(
                 )
             )
         # TODO Kirill, remove False when release edge recommendation task
-        elif task_type == "edge" and False:
+        elif task_type == Task.LINK_PREDICTION:
             loader = cast(
                 Iterable,
                 LinkNeighborLoader(
@@ -976,8 +975,10 @@ def train_1_step(
                     batch_size=self.batch, shuffle=True
                 )
             )
+        elif task_type == Task.NODE_REGRESSION:
+            raise NotImplementedError
         else:
-            raise ValueError("Unsupported task type")
+            raise ValueError(f"Unsupported task type {task_type}")
         loss = 0
         for batch in loader:
             self.before_batch(batch)
@@ -1623,6 +1624,7 @@ def train_on_batch(
             batch,
             task_type: str = None
     ) -> torch.Tensor:
+        # FIXME misha it is not task type, change to getting dvc field task
         if task_type == "multiple-graphs":
             self.optimizer.zero_grad()
             logits = self.gnn(batch.x, batch.edge_index, batch.batch)