Repository: wwu-mmll/photonai_graph (language: python)

Issue:
Seperated GraphIO from GraphConversions.
Added Documentation
GraphUtilities: Add an option for networkx graphs to be used as a list (conversion problem)



Discussion / hints:

For what use case exactly is this necessary?
This is needed for the case where networkx graphs are directly imported into the photonai-pipeline. Here photonai turns a list of networkx graphs into a numpy array. When this happens the individual networkx graphs are turned into numpy objects that are no longer networkx graphs. They become useless this way

Fix patch:
diff --git a/.github/workflows/update_documentation.yml b/.github/workflows/update_documentation.yml
index 6b94667..8866f7b 100644
--- a/.github/workflows/update_documentation.yml
+++ b/.github/workflows/update_documentation.yml
@@ -20,7 +20,7 @@ jobs:
           python-version: '3.9'
 
       - name: Install mkdocs requirements
-        run: pip install mkdocs mkdocs-material pymdown-extensions mkdocstrings-python mkdocs-git-revision-date-localized-plugin
+        run: pip install mkdocs mkdocs-material pymdown-extensions mkdocstrings-python mkdocs-jupyter
 
       - name: Install photonai_graph requirements
         run: pip install -r requirements.txt
diff --git a/documentation/README.md b/documentation/README.md
index 98ece95..bf5d3e7 100644
--- a/documentation/README.md
+++ b/documentation/README.md
@@ -9,7 +9,7 @@ To make life easier, you can simply try and install everything inside your curre
 Install mkdocs (Python package) through PyPI (inside your Python environment, if you're using one). We also need to install the material theme https://squidfunk.github.io/mkdocs-material/ and some custom plugins we're using to render Latex and so on.
 
 ```bash
-pip install mkdocs mkdocs-material pymdown-extensions mkdocstrings-python
+pip install mkdocs mkdocs-material pymdown-extensions mkdocstrings-python mkdocs-jupyter
 ```
 
 ### How to use it
diff --git a/documentation/docs/api/graph_convnets.md b/documentation/docs/api/graph_convnets.md
index 4675d92..949f39d 100644
--- a/documentation/docs/api/graph_convnets.md
+++ b/documentation/docs/api/graph_convnets.md
@@ -11,40 +11,48 @@ The Graph Neural Network module consists of three parts. The Layer Module, where
 You can also write your own custom graph neural network architecture, and register them via the PHOTON register function (link here). When writing your own custom neural nets you are free to choose your own package, as long as they implement fit, transform and predict functions like the GraphConvNet module classes. These can also be used as a blueprint if you want to integrate your own graph neural network architectures into PHOTONAI.
 
 ## DglModel
-::: photonai_graph.NeuralNets.dgl_base.DGLmodel.__init__
+::: photonai_graph.NeuralNets.dgl_base.DGLModel.__init__
     rendering:
         show_root_toc_entry: False
 
+## DGLRegressorBaseModel
+::: photonai_graph.NeuralNets.dgl_base.DGLRegressorBaseModel.__init__
+    rendering:
+        show_root_toc_entry: False
+
+## DGLClassifierBaseModel
+::: photonai_graph.NeuralNets.dgl_base.DGLClassifierBaseModel.__init__
+
 ## GCNClassifierModel
-::: photonai_graph.NeuralNets.NNClassifier.gcn_classifier.GCNClassifierModel.__init__
+::: photonai_graph.NeuralNets.GCNModel.GCNClassifierModel.__init__
     rendering:
         show_root_toc_entry: False
 
 
 ## SGConvClassifierModel
-::: photonai_graph.NeuralNets.NNClassifier.sgc_classifier.SGConvClassifierModel.__init__
+::: photonai_graph.NeuralNets.SGCModel.SGConvClassifierModel.__init__
     rendering:
         show_root_toc_entry: False
 
 
 ## GATClassifierModel
 
-::: photonai_graph.NeuralNets.NNClassifier.gat_classifier.GATClassifierModel.__init__
+::: photonai_graph.NeuralNets.GATModel.GATClassifierModel.__init__
     rendering:
         show_root_toc_entry: False
 
 
 ## GCNRegressorModel
-::: photonai_graph.NeuralNets.NNRegressor.gcn_regressor.GCNRegressorModel.__init__
+::: photonai_graph.NeuralNets.GCNModel.GCNRegressorModel.__init__
     rendering:
         show_root_toc_entry: False
 
 ## SGConvRegressorModel
-::: photonai_graph.NeuralNets.NNRegressor.sgc_regressor.SGConvRegressorModel.__init__
+::: photonai_graph.NeuralNets.SGCModel.SGConvRegressorModel.__init__
     rendering:
         show_root_toc_entry: False
 
 ## GATRegressorModel
-::: photonai_graph.NeuralNets.NNRegressor.gat_regressor.GATRegressorModel.__init__
+::: photonai_graph.NeuralNets.GATModel.GATRegressorModel.__init__
     rendering:
         show_root_toc_entry: False
\ No newline at end of file
diff --git a/documentation/docs/api/graph_io.md b/documentation/docs/api/graph_io.md
new file mode 100644
index 0000000..9dca489
--- /dev/null
+++ b/documentation/docs/api/graph_io.md
@@ -0,0 +1,7 @@
+# Graph IO functions
+
+The following Graph IO functions are available.
+
+## Graph Conversion Functions
+
+::: photonai_graph.GraphIO
diff --git a/documentation/docs/api/population_averaging.md b/documentation/docs/api/population_averaging.md
index 1ace455..08f9273 100644
--- a/documentation/docs/api/population_averaging.md
+++ b/documentation/docs/api/population_averaging.md
@@ -4,7 +4,7 @@ graphs and use the average adjacency with the individual connectivity as feature
 
 This transformer is intended for this case.
 
-## Population Averaging
+## Population Averaging Transform
 ::: photonai_graph.PopulationAveragingTransform.PopulationAveragingTransform.__init__
     rendering:
         show_root_toc_entry: False
\ No newline at end of file
diff --git a/documentation/docs/examples/GraphEmbedding_lem_example.ipynb b/documentation/docs/examples/GraphEmbedding_lem_example.ipynb
new file mode 120000
index 0000000..5ffec2a
--- /dev/null
+++ b/documentation/docs/examples/GraphEmbedding_lem_example.ipynb
@@ -0,0 +1,1 @@
+../../../examples/embeding_examples/GraphEmbedding_lem_example.ipynb
\ No newline at end of file
diff --git a/documentation/docs/examples/networkx_pipeline.ipynb b/documentation/docs/examples/networkx_pipeline.ipynb
new file mode 120000
index 0000000..77dd277
--- /dev/null
+++ b/documentation/docs/examples/networkx_pipeline.ipynb
@@ -0,0 +1,1 @@
+../../../examples/networkx_example/networkx_pipeline.ipynb
\ No newline at end of file
diff --git a/documentation/docs/load_and_save.md b/documentation/docs/load_and_save.md
index a8309e1..9e02fb3 100644
--- a/documentation/docs/load_and_save.md
+++ b/documentation/docs/load_and_save.md
@@ -12,10 +12,10 @@ From this networkx graph you can use different conversion functions to convert t
 desired format.
 
 ### Load networkx graph from file
-::: photonai_graph.GraphConversions.load_file_to_networkx
+::: photonai_graph.GraphIO.load_file_to_networkx
 
 ## Saving graphs
 Graphs can be saved as networkx 
 
 ### Graph to networkx file
-::: photonai_graph.GraphConversions.save_networkx_to_file
\ No newline at end of file
+::: photonai_graph.GraphIO.save_networkx_to_file
\ No newline at end of file
diff --git a/documentation/mkdocs.yml b/documentation/mkdocs.yml
index 93bbebe..7ea216a 100644
--- a/documentation/mkdocs.yml
+++ b/documentation/mkdocs.yml
@@ -10,7 +10,7 @@ edit_uri: 'edit/master/documentation/docs'
 extra:
   social:
     - icon: fontawesome/brands/twitter
-      link: 'https://twitter.com/medicalmachinelearninglab'
+      link: 'https://twitter.com/wwu_mmll'
 
 # Theme configuration
 theme: 
@@ -59,6 +59,7 @@ nav:
       - Examples:
          'api/examples/graph_construction_examples.md'
   - Graph Conversions: 'api/graph_conversions.md'
+  - Graph IO: 'api/graph_io.md'
   - Population Averaging: 'api/population_averaging.md'
   - Dynamic Utilities: 'api/dynamic_utils.md'
   - Graph Measures: 'api/graph_measures.md'
@@ -66,6 +67,9 @@ nav:
   - Graph Kernels: 'api/graph_kernels.md'
   - Graph Neural Networks: 'api/graph_convnets.md'
   - Graph Utilities: 'api/graph_utilities.md'
+- Examples:
+    - Networkx Loading: 'examples/networkx_pipeline.ipynb'
+    - Graph Embedding LEM: 'examples/GraphEmbedding_lem_example.ipynb'
 
 extra_css:
   - assets/stylesheets/material.css
@@ -77,4 +81,5 @@ plugins:
           selection:
             docstring_style: numpy
           rendering:
-            show_source: false
\ No newline at end of file
+            show_source: false
+  - mkdocs-jupyter
\ No newline at end of file
diff --git a/examples/LitmusTest/Networkx_Dataset.py b/examples/LitmusTest/Networkx_Dataset.py
index f1e9794..4978075 100644
--- a/examples/LitmusTest/Networkx_Dataset.py
+++ b/examples/LitmusTest/Networkx_Dataset.py
@@ -23,6 +23,7 @@ def get_sw_graph(num_graphs=500, num_nodes=20, knn=5, edge_rew=0.2):
 
     return graph_list
 
+
 def plot_nx_edge_count(graphs1, graphs2, label1="graphs_1", label2="graphs_2"):
     # plot the distribution of edges for 2 sets of graphs
     edges1 = []
diff --git a/examples/embeding_examples/GraphEmbedding_lem_example.ipynb b/examples/embeding_examples/GraphEmbedding_lem_example.ipynb
new file mode 100644
index 0000000..2823425
--- /dev/null
+++ b/examples/embeding_examples/GraphEmbedding_lem_example.ipynb
@@ -0,0 +1,409 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": true
+   },
+   "source": [
+    "# Graph Embeding LEM Example\n",
+    "\n",
+    "This example will show you how to use the LEM graph embedding with a PHOTONAI Graph Pipeline\n",
+    "\n",
+    "At first we have to import the requirements"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from photonai.base import Hyperpipe, PipelineElement\n",
+    "from photonai_graph.GraphUtilities import get_random_connectivity_data, get_random_labels\n",
+    "from sklearn.model_selection import KFold"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we are able to generate some sample data for a classification"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# make random matrices to simulate connectivity matrices\n",
+    "X = get_random_connectivity_data(number_of_nodes=50, number_of_individuals=100)\n",
+    "y = get_random_labels(l_type=\"classification\", number_of_labels=100)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The next step is the definition of our PHOTONAI pipeline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using backend: pytorch\n"
+     ]
+    }
+   ],
+   "source": [
+    "my_pipe = Hyperpipe('basic_gembedding_pipe',\n",
+    "                    inner_cv=KFold(n_splits=5),\n",
+    "                    outer_cv=KFold(n_splits=5),\n",
+    "                    optimizer='sk_opt',\n",
+    "                    optimizer_params={'n_configurations': 25},\n",
+    "                    metrics=['accuracy', 'balanced_accuracy', 'recall', 'precision'],\n",
+    "                    best_config_metric='accuracy')\n",
+    "\n",
+    "my_pipe.add(PipelineElement('GraphConstructorThreshold', threshold=0.95))\n",
+    "\n",
+    "my_pipe.add(PipelineElement('GraphEmbeddingLocallyLinearEmbedding'))\n",
+    "\n",
+    "my_pipe.add(PipelineElement('SVC'))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Finally we can fit the pipeline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "=====================================================================================================\n",
+      "PHOTONAI ANALYSIS: basic_gembedding_pipe\n",
+      "=====================================================================================================\n",
+      "\n",
+      "*****************************************************************************************************\n",
+      "Outer Cross validation Fold 1\n",
+      "*****************************************************************************************************\n",
+      "Did not find any hyperparameter to convert into skopt space.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/photonai/optimization/scikit_optimize/sk_opt.py:169: UserWarning: Did not find any hyperparameter to convert into skopt space.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "BEST_CONFIG \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "{}\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "VALIDATION PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7188      |      0.4625      |\n",
+      "| balanced_accuracy |       0.7065      |      0.4510      |\n",
+      "|       recall      |       0.8645      |      0.7548      |\n",
+      "|     precision     |       0.6906      |      0.5097      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "TEST PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.6750      |      0.4500      |\n",
+      "| balanced_accuracy |       0.6637      |      0.5208      |\n",
+      "|       recall      |       0.8140      |      0.8750      |\n",
+      "|     precision     |       0.6604      |      0.4118      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "\n",
+      "*****************************************************************************************************\n",
+      "Outer Cross validation Fold 2\n",
+      "*****************************************************************************************************\n",
+      "Did not find any hyperparameter to convert into skopt space.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/photonai/optimization/scikit_optimize/sk_opt.py:169: UserWarning: Did not find any hyperparameter to convert into skopt space.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "BEST_CONFIG \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "{}\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "VALIDATION PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7000      |      0.4375      |\n",
+      "| balanced_accuracy |       0.6906      |      0.5035      |\n",
+      "|       recall      |       0.5947      |      0.4344      |\n",
+      "|     precision     |       0.7426      |      0.4583      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "TEST PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.6750      |      0.4500      |\n",
+      "| balanced_accuracy |       0.6723      |      0.4583      |\n",
+      "|       recall      |       0.5641      |      0.4167      |\n",
+      "|     precision     |       0.7097      |      0.5556      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "\n",
+      "*****************************************************************************************************\n",
+      "Outer Cross validation Fold 3\n",
+      "*****************************************************************************************************\n",
+      "Did not find any hyperparameter to convert into skopt space.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/photonai/optimization/scikit_optimize/sk_opt.py:169: UserWarning: Did not find any hyperparameter to convert into skopt space.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "BEST_CONFIG \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "{}\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "VALIDATION PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.6875      |      0.3750      |\n",
+      "| balanced_accuracy |       0.6778      |      0.3962      |\n",
+      "|       recall      |       0.7550      |      0.6338      |\n",
+      "|     precision     |       0.6772      |      0.4459      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "TEST PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7750      |      0.2000      |\n",
+      "| balanced_accuracy |       0.7694      |      0.2121      |\n",
+      "|       recall      |       0.8810      |      0.3333      |\n",
+      "|     precision     |       0.7400      |      0.2308      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "\n",
+      "*****************************************************************************************************\n",
+      "Outer Cross validation Fold 4\n",
+      "*****************************************************************************************************\n",
+      "Did not find any hyperparameter to convert into skopt space.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/photonai/optimization/scikit_optimize/sk_opt.py:169: UserWarning: Did not find any hyperparameter to convert into skopt space.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "BEST_CONFIG \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "{}\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "VALIDATION PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.6625      |      0.4000      |\n",
+      "| balanced_accuracy |       0.6599      |      0.4284      |\n",
+      "|       recall      |       0.6350      |      0.2922      |\n",
+      "|     precision     |       0.6640      |      0.3683      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "TEST PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7250      |      0.4000      |\n",
+      "| balanced_accuracy |       0.7236      |      0.4375      |\n",
+      "|       recall      |       0.6667      |      0.2500      |\n",
+      "|     precision     |       0.7429      |      0.5000      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "\n",
+      "*****************************************************************************************************\n",
+      "Outer Cross validation Fold 5\n",
+      "*****************************************************************************************************\n",
+      "Did not find any hyperparameter to convert into skopt space.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/photonai/optimization/scikit_optimize/sk_opt.py:169: UserWarning: Did not find any hyperparameter to convert into skopt space.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "BEST_CONFIG \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "{}\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "VALIDATION PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7188      |      0.4875      |\n",
+      "| balanced_accuracy |       0.7089      |      0.5645      |\n",
+      "|       recall      |       0.7702      |      0.6361      |\n",
+      "|     precision     |       0.7011      |      0.6088      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "TEST PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7125      |      0.3500      |\n",
+      "| balanced_accuracy |       0.7145      |      0.3500      |\n",
+      "|       recall      |       0.6341      |      0.5000      |\n",
+      "|     precision     |       0.7647      |      0.3846      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "*****************************************************************************************************\n",
+      "Finished all outer fold computations.\n",
+      "*****************************************************************************************************\n",
+      "\n",
+      "ANALYSIS INFORMATION ================================================================================ \n",
+      "Project Folder: /Users/jan/PycharmProjects/photonai_graph/examples/embeding_examples/basic_gembedding_pipe_results_2022-03-17_09-39-12,\n",
+      "Computation Time: 2022-03-17 09:39:12.844667 - 2022-03-17 09:40:11.639869\n",
+      "Duration: 0:00:58.795202\n",
+      "Optimized for: accuracy\n",
+      "Hyperparameter Optimizer: sk_opt\n",
+      "\n",
+      "DUMMY RESULTS =======================================================================================\n",
+      "+-------------------+--+\n",
+      "| PERFORMANCE DUMMY |  |\n",
+      "+-------------------+--+\n",
+      "+-------------------+--+\n",
+      "\n",
+      "AVERAGE PERFORMANCE ACROSS OUTER FOLDS ==============================================================\n",
+      "+-------------------+---------------+--------------+-----------+----------+\n",
+      "|    Metric Name    | Training Mean | Training Std | Test Mean | Test Std |\n",
+      "+-------------------+---------------+--------------+-----------+----------+\n",
+      "|      accuracy     |     0.7125    |   0.037081   |    0.37   | 0.092736 |\n",
+      "| balanced_accuracy |    0.708708   |   0.038183   |  0.395758 | 0.106881 |\n",
+      "|       recall      |    0.711964   |   0.117413   |   0.475   | 0.216667 |\n",
+      "|     precision     |    0.723524   |   0.036112   |  0.416541 | 0.111215 |\n",
+      "+-------------------+---------------+--------------+-----------+----------+\n",
+      "\n",
+      "BEST HYPERPARAMETER CONFIGURATION ===================================================================\n",
+      "{}\n",
+      "\n",
+      "\n",
+      "+--------+----------+-------------------+--------+-----------+----------------------------+\n",
+      "| fold # | accuracy | balanced_accuracy | recall | precision | Best Hyperparameter Config |\n",
+      "+--------+----------+-------------------+--------+-----------+----------------------------+\n",
+      "|   1*   |  0.4500  |       0.5208      | 0.8750 |   0.4118  |             {}             |\n",
+      "|   2    |  0.4500  |       0.4583      | 0.4167 |   0.5556  |             {}             |\n",
+      "|   3    |  0.2000  |       0.2121      | 0.3333 |   0.2308  |             {}             |\n",
+      "|   4    |  0.4000  |       0.4375      | 0.2500 |   0.5000  |             {}             |\n",
+      "|   5    |  0.3500  |       0.3500      | 0.5000 |   0.3846  |             {}             |\n",
+      "+--------+----------+-------------------+--------+-----------+----------------------------+\n",
+      "\n",
+      "PHOTONAI 2.2.0 ======================================================================================\n",
+      "Your results are stored in /Users/jan/PycharmProjects/photonai_graph/examples/embeding_examples/basic_gembedding_pipe_results_2022-03-17_09-39-12\n",
+      "Go to https://explorer.photon-ai.com and upload your photon_result_file.json for convenient result visualization! \n",
+      "For more info and documentation visit https://www.photon-ai.com\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Hyperpipe(name='basic_gembedding_pipe')"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "my_pipe.fit(X, y)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 1
+}
\ No newline at end of file
diff --git a/examples/networkx_example/networkx_pipeline.ipynb b/examples/networkx_example/networkx_pipeline.ipynb
new file mode 100644
index 0000000..dca1389
--- /dev/null
+++ b/examples/networkx_example/networkx_pipeline.ipynb
@@ -0,0 +1,292 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Networkx example\n",
+    "\n",
+    "This Example will show you step by step, how to use networkx graphs with a PHOTONAI Graph pipeline.\n",
+    "\n",
+    "First: Necessary imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "pycharm": {
+     "is_executing": true
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using backend: pytorch\n"
+     ]
+    }
+   ],
+   "source": [
+    "from networkx.generators import barabasi_albert_graph, erdos_renyi_graph\n",
+    "from photonai.base import Hyperpipe, PipelineElement\n",
+    "from sklearn.model_selection import KFold\n",
+    "import numpy as np\n",
+    "\n",
+    "from photonai_graph.GraphConversions import convert_graphs\n",
+    "from photonai_graph.GraphConversions import networkx_to_dense"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next: Generate graphs and labels"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# create dataset of 1000 graphs\n",
+    "rr_graphs = [erdos_renyi_graph(20, .25) for _ in range(500)]\n",
+    "er_graphs = [erdos_renyi_graph(20, .2) for _ in range(500)]\n",
+    "graphs = rr_graphs + er_graphs\n",
+    "\n",
+    "# we have to transform the networkx graphs into numpy graphs before passing them to photon\n",
+    "graphs = np.array(convert_graphs(graphs, output_format=\"dense\"))\n",
+    "graphs = np.expand_dims(graphs, axis=-1)\n",
+    "\n",
+    "# create labels\n",
+    "rr_labels = [0] * 500\n",
+    "er_labels = [1] * 500\n",
+    "labels = rr_labels + er_labels"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup Hyperpipe"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Design your Pipeline\n",
+    "my_pipe = Hyperpipe('networkx_example_pipe',\n",
+    "                    inner_cv=KFold(n_splits=2),\n",
+    "                    optimizer='sk_opt',\n",
+    "                    optimizer_params={'n_configurations': 25},\n",
+    "                    metrics=['accuracy', 'balanced_accuracy', 'recall', 'precision'],\n",
+    "                    best_config_metric='accuracy')\n",
+    "\n",
+    "my_pipe.add(PipelineElement('GCNClassifier', feature_axis=0))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Final step: Train your Hyperpipe"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "=====================================================================================================\n",
+      "PHOTONAI ANALYSIS: networkx_example_pipe\n",
+      "=====================================================================================================\n",
+      "\n",
+      "*****************************************************************************************************\n",
+      "Outer Cross validation Fold 1\n",
+      "*****************************************************************************************************\n",
+      "Did not find any hyperparameter to convert into skopt space.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/photonai/optimization/scikit_optimize/sk_opt.py:169: UserWarning: Did not find any hyperparameter to convert into skopt space.\n",
+      "  warnings.warn(msg)\n",
+      "  0%|          | 0/200 [00:00<?, ?it/s]/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
+      "  warnings.warn('User provided device_type of \\'cuda\\', but CUDA is not available. Disabling')\n",
+      "100%|██████████| 200/200 [00:22<00:00,  8.91it/s]\n",
+      "  0%|          | 0/200 [00:00<?, ?it/s]/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
+      "  warnings.warn('User provided device_type of \\'cuda\\', but CUDA is not available. Disabling')\n",
+      "100%|██████████| 200/200 [00:21<00:00,  9.17it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "BEST_CONFIG \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "{}\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "VALIDATION PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7875      |      0.7712      |\n",
+      "| balanced_accuracy |       0.7864      |      0.7715      |\n",
+      "|       recall      |       0.7680      |      0.7497      |\n",
+      "|     precision     |       0.8035      |      0.7864      |\n",
+      "+-------------------+-------------------+------------------+\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 0/200 [00:00<?, ?it/s]/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
+      "  warnings.warn('User provided device_type of \\'cuda\\', but CUDA is not available. Disabling')\n",
+      "100%|██████████| 200/200 [00:41<00:00,  4.77it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "-----------------------------------------------------------------------------------------------------\n",
+      "TEST PERFORMANCE\n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|       METRIC      | PERFORMANCE TRAIN | PERFORMANCE TEST |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "|      accuracy     |       0.7788      |      0.7750      |\n",
+      "| balanced_accuracy |       0.7787      |      0.7737      |\n",
+      "|       recall      |       0.7816      |      0.7320      |\n",
+      "|     precision     |       0.7797      |      0.7889      |\n",
+      "+-------------------+-------------------+------------------+\n",
+      "*****************************************************************************************************\n",
+      "Finished all outer fold computations.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 0/200 [00:00<?, ?it/s]/Users/jan/anaconda3/envs/photonai_graph_test/lib/python3.9/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
+      "  warnings.warn('User provided device_type of \\'cuda\\', but CUDA is not available. Disabling')\n",
+      "100%|██████████| 200/200 [00:54<00:00,  3.66it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "*****************************************************************************************************\n",
+      "\n",
+      "ANALYSIS INFORMATION ================================================================================ \n",
+      "Project Folder: /Users/jan/PycharmProjects/photonai_graph/examples/networkx_example/networkx_example_pipe_results_2022-03-17_09-31-35,\n",
+      "Computation Time: 2022-03-17 09:31:35.916935 - 2022-03-17 09:33:04.107795\n",
+      "Duration: 0:01:28.190860\n",
+      "Optimized for: accuracy\n",
+      "Hyperparameter Optimizer: sk_opt\n",
+      "\n",
+      "DUMMY RESULTS =======================================================================================\n",
+      "+-------------------+--+\n",
+      "| PERFORMANCE DUMMY |  |\n",
+      "+-------------------+--+\n",
+      "+-------------------+--+\n",
+      "\n",
+      "AVERAGE PERFORMANCE ACROSS OUTER FOLDS ==============================================================\n",
+      "+-------------------+---------------+--------------+-----------+----------+\n",
+      "|    Metric Name    | Training Mean | Training Std | Test Mean | Test Std |\n",
+      "+-------------------+---------------+--------------+-----------+----------+\n",
+      "|      accuracy     |    0.77875    |     0.0      |   0.775   |   0.0    |\n",
+      "| balanced_accuracy |    0.778728   |     0.0      |  0.773746 |   0.0    |\n",
+      "|       recall      |    0.781638   |     0.0      |  0.731959 |   0.0    |\n",
+      "|     precision     |    0.779703   |     0.0      |  0.788889 |   0.0    |\n",
+      "+-------------------+---------------+--------------+-----------+----------+\n",
+      "\n",
+      "BEST HYPERPARAMETER CONFIGURATION ===================================================================\n",
+      "{}\n",
+      "\n",
+      "\n",
+      "+--------+----------+-------------------+--------+-----------+----------------------------+\n",
+      "| fold # | accuracy | balanced_accuracy | recall | precision | Best Hyperparameter Config |\n",
+      "+--------+----------+-------------------+--------+-----------+----------------------------+\n",
+      "|   1*   |  0.7750  |       0.7737      | 0.7320 |   0.7889  |             {}             |\n",
+      "+--------+----------+-------------------+--------+-----------+----------------------------+\n",
+      "\n",
+      "PHOTONAI 2.2.0 ======================================================================================\n",
+      "Your results are stored in /Users/jan/PycharmProjects/photonai_graph/examples/networkx_example/networkx_example_pipe_results_2022-03-17_09-31-35\n",
+      "Go to https://explorer.photon-ai.com and upload your photon_result_file.json for convenient result visualization! \n",
+      "For more info and documentation visit https://www.photon-ai.com\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Hyperpipe(name='networkx_example_pipe')"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "my_pipe.fit(graphs, labels)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 1
+}
\ No newline at end of file
diff --git a/photonai_graph/GraphConversions.py b/photonai_graph/GraphConversions.py
index 639eeab..f2ba010 100644
--- a/photonai_graph/GraphConversions.py
+++ b/photonai_graph/GraphConversions.py
@@ -1,49 +1,15 @@
-from networkx.drawing.nx_pydot import write_dot, read_dot
 from scipy import sparse
 import networkx as nx
 import numpy as np
-import os
 import warnings
-import yaml
 from photonai_graph.util import assert_imported
 
 try:
     import dgl
+    import torch
 except ImportError:
     pass
 
-output_formats = {
-    "dot": write_dot,
-    "AdjacencyList": nx.write_adjlist,
-    "MultilineAdjacencyList": nx.write_multiline_adjlist,
-    "EdgeList": nx.write_edgelist,
-    "WeightedEdgeList": nx.write_weighted_edgelist,
-    "GEXF": nx.write_gexf,
-    "pickle": nx.write_gpickle,
-    "GLM": nx.write_gml,
-    "GraphML": nx.write_graphml,
-    "GraphML-XML": nx.write_graphml_xml,
-    "GraphML-LXML": nx.write_graphml_lxml,
-    "YAML": yaml.dump,
-    "graph6": nx.write_graph6,
-    "PAJEK": nx.write_pajek
-}
-
-input_formats = {
-    "dot": read_dot,
-    "AdjacencyList": nx.read_adjlist,
-    "MultilineAdjacencyList": nx.read_multiline_adjlist,
-    "EdgeList": nx.read_edgelist,
-    "WeightedEdgeList": nx.read_edgelist,
-    "GEXF": nx.read_gexf,
-    "pickle": nx.read_gpickle,
-    "GML": nx.read_gml,
-    "GraphML": nx.read_graphml,
-    "YAML": yaml.load,
-    "graph6": nx.read_graph6,
-    "PAJEK": nx.read_pajek
-}
-
 sparse_types = {
     "bsr_matrix": sparse.bsr_matrix,
     "coo_matrix": sparse.coo_matrix,
@@ -56,90 +22,6 @@
 }
 
 
-def save_networkx_to_file(graphs, path, output_format="dot", ids=None):
-    """Saves networkx graph(s) to a file
-
-        Parameters
-        ----------
-        graphs: a list of networkx graphs or single networkx graph
-            list of graphs you want to save
-        path: str
-            path where to save the data as a string
-        output_format: str
-            output format in which the graphs should be saved. See output_formats
-            at top of the file for more information on valid output formats
-        ids: list, default=None
-            a list of ids after which to name each graph in the file. If no list is
-            specified the graphs are just numbered as graph_x
-
-    """
-    def save_single_networkx_to_file(graph, path, output_format="dot"):
-        """internal helper function
-
-        Parameters
-        ----------
-        graph
-            networkx graph object
-        path
-            path to save the graph to
-        output_format
-            desired output format
-        """
-        if not isinstance(graph, nx.classes.graph.Graph):
-            raise ValueError(f"Got unknown object for serialization: {type(graph)}")
-        output_formats[output_format](graph, path)
-
-    if output_format not in output_formats:
-        raise ValueError("The output format is not supported. Please check your output format.")
-
-    if not isinstance(graphs, list):
-        return save_single_networkx_to_file(graph=graphs, path=path, output_format=output_format)
-
-    # Check if we have got a list of ids
-    if ids is None or not isinstance(ids, list):
-        ids = np.arange(len(graphs))
-
-    # check if id and graphs they have equal length
-    if len(graphs) != len(ids):
-        raise ValueError(
-            'The photonai_graph ID list and the list of Graphs are not of equal length. '
-            'Please ensure that they have equal length.')
-
-    for graph, i in zip(graphs, ids):
-        graph_filename = "graph_" + str(i)
-        graph_path = os.path.join(path, graph_filename)
-        save_single_networkx_to_file(graph=graph, path=graph_path, output_format=output_format)
-
-
-def load_file_to_networkx(path, input_format="dot"):
-    """load graph into networkx from a file
-
-        Parameters
-        ----------
-        path: list, str
-            path(s) where graphs are stored
-        input_format:
-            format in which these graphs are stored
-
-        Returns
-        -------
-        list
-            List of loaded networkx graphs
-    """
-    if isinstance(path, str):
-        path = [path]
-
-    graph_list = []
-    if input_format not in input_formats:
-        raise KeyError("Input format is not supported right now.")
-
-    for graph in path:
-        g = input_formats[input_format](graph)
-        graph_list.append(g)
-
-    return graph_list
-
-
 def networkx_to_dense(graphs):
     """convert a networkx graph to a numpy array
 
@@ -457,9 +339,16 @@ def dense_to_dgl(graphs, adjacency_axis=None, feature_axis=None):
 
     if adjacency_axis is None:
         raise NotImplementedError('dense to dgl not implemented without adjacency axis')
-    else:
-        sparse_mtrx = dense_to_sparse(graphs, adjacency_axis=adjacency_axis, feature_axis=feature_axis)
-        graph_list = sparse_to_dgl(sparse_mtrx)
+    if not isinstance(graphs, np.ndarray):
+        raise ValueError(f"Expected np.ndarray as input, got {type(graphs)}")
+
+    graph_list = []
+    for graph in range(graphs.shape[0]):
+        src, dst = np.nonzero(graphs[graph, :, :, adjacency_axis])
+        g = dgl.graph((src, dst), num_nodes=graphs[graph, ...].shape[1])
+        feat = torch.tensor(graphs[graph, :, :, feature_axis])
+        g.ndata['feat'] = feat
+        graph_list.append(g)
     return graph_list
 
 
diff --git a/photonai_graph/GraphIO.py b/photonai_graph/GraphIO.py
new file mode 100644
index 0000000..326be4a
--- /dev/null
+++ b/photonai_graph/GraphIO.py
@@ -0,0 +1,122 @@
+import os
+
+from networkx.drawing.nx_pydot import write_dot, read_dot
+import networkx as nx
+import numpy as np
+import yaml
+
+output_formats = {
+    "dot": write_dot,
+    "AdjacencyList": nx.write_adjlist,
+    "MultilineAdjacencyList": nx.write_multiline_adjlist,
+    "EdgeList": nx.write_edgelist,
+    "WeightedEdgeList": nx.write_weighted_edgelist,
+    "GEXF": nx.write_gexf,
+    "pickle": nx.write_gpickle,
+    "GLM": nx.write_gml,
+    "GraphML": nx.write_graphml,
+    "GraphML-XML": nx.write_graphml_xml,
+    "GraphML-LXML": nx.write_graphml_lxml,
+    "YAML": yaml.dump,
+    "graph6": nx.write_graph6,
+    "PAJEK": nx.write_pajek
+}
+
+input_formats = {
+    "dot": read_dot,
+    "AdjacencyList": nx.read_adjlist,
+    "MultilineAdjacencyList": nx.read_multiline_adjlist,
+    "EdgeList": nx.read_edgelist,
+    "WeightedEdgeList": nx.read_edgelist,
+    "GEXF": nx.read_gexf,
+    "pickle": nx.read_gpickle,
+    "GML": nx.read_gml,
+    "GraphML": nx.read_graphml,
+    "YAML": yaml.load,
+    "graph6": nx.read_graph6,
+    "PAJEK": nx.read_pajek
+}
+
+
+def save_networkx_to_file(graphs, path, output_format="dot", ids=None):
+    """Saves networkx graph(s) to a file
+
+        Parameters
+        ----------
+        graphs: a list of networkx graphs or single networkx graph
+            list of graphs you want to save
+        path: str
+            path where to save the data as a string
+        output_format: str
+            output format in which the graphs should be saved. See output_formats
+            at top of the file for more information on valid output formats
+        ids: list, default=None
+            a list of ids after which to name each graph in the file. If no list is
+            specified the graphs are just numbered as graph_x
+
+    """
+    def save_single_networkx_to_file(graph, path, output_format="dot"):
+        """internal helper function
+
+        Parameters
+        ----------
+        graph
+            networkx graph object
+        path
+            path to save the graph to
+        output_format
+            desired output format
+        """
+        if not isinstance(graph, nx.classes.graph.Graph):
+            raise ValueError(f"Got unknown object for serialization: {type(graph)}")
+        output_formats[output_format](graph, path)
+
+    if output_format not in output_formats:
+        raise ValueError("The output format is not supported. Please check your output format.")
+
+    if not isinstance(graphs, list):
+        return save_single_networkx_to_file(graph=graphs, path=path, output_format=output_format)
+
+    # Check if we have got a list of ids
+    if ids is None or not isinstance(ids, list):
+        ids = np.arange(len(graphs))
+
+    # check if id and graphs they have equal length
+    if len(graphs) != len(ids):
+        raise ValueError(
+            'The photonai_graph ID list and the list of Graphs are not of equal length. '
+            'Please ensure that they have equal length.')
+
+    for graph, i in zip(graphs, ids):
+        graph_filename = "graph_" + str(i)
+        graph_path = os.path.join(path, graph_filename)
+        save_single_networkx_to_file(graph=graph, path=graph_path, output_format=output_format)
+
+
+def load_file_to_networkx(path, input_format="dot"):
+    """load graph into networkx from a file
+
+        Parameters
+        ----------
+        path: list, str
+            path(s) where graphs are stored
+        input_format:
+            format in which these graphs are stored
+
+        Returns
+        -------
+        list
+            List of loaded networkx graphs
+    """
+    if isinstance(path, str):
+        path = [path]
+
+    graph_list = []
+    if input_format not in input_formats:
+        raise KeyError("Input format is not supported right now.")
+
+    for graph in path:
+        g = input_formats[input_format](graph)
+        graph_list.append(g)
+
+    return graph_list
diff --git a/photonai_graph/GraphUtilities.py b/photonai_graph/GraphUtilities.py
index 127945f..0afd00c 100644
--- a/photonai_graph/GraphUtilities.py
+++ b/photonai_graph/GraphUtilities.py
@@ -8,7 +8,6 @@
 from scipy import sparse
 import os
 import matplotlib.pyplot as plt
-from photonai_graph.GraphConversions import save_networkx_to_file
 
 
 def draw_connectogram(graph, edge_rad=None, colorscheme=None, nodesize=None,
@@ -293,41 +292,6 @@ def get_random_labels(l_type="classification", number_of_labels=10):
     return y
 
 
-def save_graphs(graphs, path="", input_format="networkx", output_format="dot", ids=None):
-    """save graphs to file.
-
-        Parameters
-        ----------
-        graphs :
-            a list or a np.ndarray of graphs to be saved
-
-        input_format : str, default="networkx"
-            format of the graphs to be saved
-
-        path : str, default=""
-            path where to save the graphs
-
-        output_format : str, default="dot"
-            the output format in which to save the graphs
-
-        ids :
-            a list containing the ids of the graphs. Must have same length as the graph list or np.ndarray.
-
-
-        Examples
-        --------
-        >>> g1, g2 = nx.line_graph(10), nx.line_graph(7)
-        >>> graphs = [g1, g2]
-        >>> save_graphs(graphs, path="path/to/your/data/")
-
-        """
-    # check input format
-    if input_format == "networkx":
-        save_networkx_to_file(graphs, path, output_format=output_format, ids=ids)
-    else:
-        raise NotImplementedError("Your desired output format is not supported yet.")
-
-
 def visualize_networkx(graphs, layout=nx.spring_layout, colorscheme="Blues", show=True):
     """Visualize a networkx graph or graphs using networkx built-in visualization.
 
diff --git a/photonai_graph/NeuralNets/GATModel.py b/photonai_graph/NeuralNets/GATModel.py
new file mode 100644
index 0000000..68c55d6
--- /dev/null
+++ b/photonai_graph/NeuralNets/GATModel.py
@@ -0,0 +1,157 @@
+from typing import List
+
+import dgl
+import torch.nn as nn
+import numpy as np
+from dgl.nn.pytorch import GATConv
+
+from photonai_graph.NeuralNets.dgl_base import DGLClassifierBaseModel, DGLRegressorBaseModel
+
+
+class GATModel(nn.Module):
+    def __init__(self, in_dim, hidden_dim, num_heads, n_classes, hidden_layers,
+                 agg_mode, allow_zero_in_degree):
+        super(GATModel, self).__init__()
+
+        self.agg_mode = agg_mode
+        self.layers = nn.ModuleList()
+        # append gat layer according to inputs
+        self.layers.append(GATConv(in_dim, hidden_dim, num_heads[0],
+                                   allow_zero_in_degree=allow_zero_in_degree))
+        # hidden layers
+        for layer in range(1, hidden_layers):
+            self.layers.append(GATConv(
+                hidden_dim, hidden_dim, num_heads[layer],
+                allow_zero_in_degree=allow_zero_in_degree))
+        # output layer
+        self.classify = nn.Linear(hidden_dim, n_classes)
+
+    def forward(self, bg):
+        # For undirected graphs, in_degree is the same as
+        # out_degree.
+        h = bg.in_degrees().view(-1, 1).float()
+        for i, gnn in enumerate(self.layers):
+            h = gnn(bg, h)
+            if self.agg_mode == 'flatten':
+                h = h.flatten(1)
+            else:
+                h = h.mean(1)
+
+        bg.ndata['h'] = h
+        hg = dgl.mean_nodes(bg, 'h')
+        return self.classify(hg)
+
+
+class GATClassifierModel(DGLClassifierBaseModel):
+
+    def __init__(self,
+                 in_dim: int = 1,
+                 hidden_layers: int = 2,
+                 hidden_dim: int = 256,
+                 heads: List = None,
+                 agg_mode="mean",
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = ''):
+        """
+            Graph Attention Network for graph classification. GAT Layers
+            are modeled after Veličković et al., 2018.
+            Implementation based on dgl & pytorch.
+
+
+            Parameters
+            ----------
+            in_dim: int,default=1
+                input dimension
+            hidden_layers: int,default=2
+                number of hidden layers used by the model
+            hidden_dim: int,default=256
+                dimensions in the hidden layers
+            heads: list,default=None
+                list with number of heads per hidden layer
+            agg_mode: str, default='mean'
+                aggregation mode for the graph convolutional layers
+
+        """
+        super(GATClassifierModel, self).__init__(nn_epochs=nn_epochs,
+                                                 learning_rate=learning_rate,
+                                                 batch_size=batch_size,
+                                                 adjacency_axis=adjacency_axis,
+                                                 feature_axis=feature_axis,
+                                                 add_self_loops=add_self_loops,
+                                                 allow_zero_in_degree=allow_zero_in_degree,
+                                                 logs=logs)
+        if heads is None:
+            heads = [2, 2]
+            # Todo: if heads is not length of hidden layers +1 (bc of the first layer)
+        self.in_dim = in_dim
+        self.hidden_layers = hidden_layers
+        self.hidden_dim = hidden_dim
+        self.heads = heads
+        self.agg_mode = agg_mode
+
+    def _init_model(self, X=None, y=None):
+        self.model = GATModel(self.in_dim, self.hidden_dim, self.heads,
+                              len(np.unique(y)), self.hidden_layers, self.agg_mode,
+                              allow_zero_in_degree=self.allow_zero_in_degree)
+
+
+class GATRegressorModel(DGLRegressorBaseModel):
+
+    def __init__(self,
+                 in_dim: int = 1,
+                 hidden_layers: int = 2,
+                 hidden_dim: int = 256,
+                 heads: List = None,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = None,
+                 agg_mode: str = None):
+        """
+            Graph Attention Network for graph regression. GAT Layers
+            are modeled after Veličković et al., 2018. Implementation
+            based on dgl & pytorch.
+
+
+            Parameters
+            ----------
+            in_dim: int,default=1
+                input dimension
+            hidden_layers: int,default=2
+                number of hidden layers used by the model
+            hidden_dim: int,default=256
+                dimensions in the hidden layers
+            heads: list,default=None
+                list with number of heads per hidden layer
+
+        """
+        super(GATRegressorModel, self).__init__(nn_epochs=nn_epochs,
+                                                learning_rate=learning_rate,
+                                                batch_size=batch_size,
+                                                adjacency_axis=adjacency_axis,
+                                                feature_axis=feature_axis,
+                                                add_self_loops=add_self_loops,
+                                                allow_zero_in_degree=allow_zero_in_degree,
+                                                logs=logs)
+        if heads is None:
+            heads = [2, 2]
+            # Todo: if heads is not length of hidden layers +1 (bc of the first layer)
+        self.in_dim = in_dim
+        self.hidden_layers = hidden_layers
+        self.hidden_dim = hidden_dim
+        self.heads = heads
+        self.agg_mode = agg_mode
+
+    def _init_model(self, X=None, y=None):
+        self.model = GATModel(self.in_dim, self.hidden_dim, self.heads, 1, self.hidden_layers,
+                              allow_zero_in_degree=self.allow_zero_in_degree, agg_mode=self.agg_mode).float()
diff --git a/photonai_graph/NeuralNets/GCNModel.py b/photonai_graph/NeuralNets/GCNModel.py
new file mode 100644
index 0000000..b2b24f7
--- /dev/null
+++ b/photonai_graph/NeuralNets/GCNModel.py
@@ -0,0 +1,129 @@
+import dgl
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.nn.pytorch import GraphConv
+
+from photonai_graph.NeuralNets.dgl_base import DGLRegressorBaseModel, DGLClassifierBaseModel
+
+
+class GCNClassifier(nn.Module):
+    def __init__(self, in_dim, hidden_dim, n_classes, hidden_layers, allow_zero_in_degree):
+        super(GCNClassifier, self).__init__()
+        self.layers = nn.ModuleList()
+        # input layers
+        self.layers.append(GraphConv(in_dim, hidden_dim, allow_zero_in_degree=allow_zero_in_degree))
+        # hidden layers
+        for layer in range(1, hidden_layers):
+            self.layers.append(GraphConv(hidden_dim, hidden_dim, allow_zero_in_degree=allow_zero_in_degree))
+        self.classify = nn.Linear(hidden_dim, n_classes)
+
+    def forward(self, g):
+        # Use node degree as the initial node feature. For undirected graphs, the in-degree
+        # is the same as the out_degree.
+        h = g.in_degrees().view(-1, 1).float()
+        # Perform graph convolution and activation function.
+        for i, gnn in enumerate(self.layers):
+            h = F.relu(gnn(g, h))
+        g.ndata['h'] = h
+        # Calculate graph representation by averaging all the node representations.
+        hg = dgl.mean_nodes(g, 'h')
+        return self.classify(hg)
+
+
+class GCNClassifierModel(DGLClassifierBaseModel):
+
+    def __init__(self,
+                 in_dim: int = 1,
+                 hidden_layers: int = 2,
+                 hidden_dim: int = 256,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = ''):
+        """
+        Graph Attention Network for graph classification. GCN Layers
+        from Kipf & Welling, 2017.
+        Implementation based on dgl & pytorch.
+
+
+        Parameters
+        ----------
+        in_dim: int,default=1
+            input dimension
+        hidden_layers: int,default=2
+            number of hidden layers used by the model
+        hidden_dim: int,default=256
+            dimensions in the hidden layers
+
+        """
+        super(GCNClassifierModel, self).__init__(nn_epochs=nn_epochs,
+                                                 learning_rate=learning_rate,
+                                                 batch_size=batch_size,
+                                                 adjacency_axis=adjacency_axis,
+                                                 feature_axis=feature_axis,
+                                                 add_self_loops=add_self_loops,
+                                                 allow_zero_in_degree=allow_zero_in_degree,
+                                                 logs=logs)
+        self.in_dim = in_dim
+        self.hidden_dim = hidden_dim
+        self.hidden_layers = hidden_layers
+
+    def _init_model(self, X=None, y=None):
+        self.model = GCNClassifier(self.in_dim,
+                                   self.hidden_dim,
+                                   len(np.unique(y)),
+                                   self.hidden_layers,
+                                   allow_zero_in_degree=self.allow_zero_in_degree)
+
+
+class GCNRegressorModel(DGLRegressorBaseModel):
+
+    def __init__(self,
+                 in_dim: int = 1,
+                 hidden_layers: int = 2,
+                 hidden_dim: int = 256,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = ''):
+        """
+        Graph convolutional Network for graph regression. GCN Layers
+        from Kipf & Welling, 2017. Implementation based on dgl & pytorch.
+
+
+        Parameters
+        ----------
+        in_dim: int,default=1
+            input dimension
+        hidden_layers: int,default=2
+            number of hidden layers used by the model
+        hidden_dim: int,default=256
+            dimensions in the hidden layers
+
+        """
+        super(GCNRegressorModel, self).__init__(nn_epochs=nn_epochs,
+                                                learning_rate=learning_rate,
+                                                batch_size=batch_size,
+                                                adjacency_axis=adjacency_axis,
+                                                feature_axis=feature_axis,
+                                                add_self_loops=add_self_loops,
+                                                allow_zero_in_degree=allow_zero_in_degree,
+                                                logs=logs)
+        self.in_dim = in_dim
+        self.hidden_dim = hidden_dim
+        self.hidden_layers = hidden_layers
+
+    def _init_model(self, X=None, y=None):
+        self.model = GCNClassifier(self.in_dim,
+                                   self.hidden_dim, 1,
+                                   self.hidden_layers,
+                                   allow_zero_in_degree=self.allow_zero_in_degree).float()
diff --git a/photonai_graph/NeuralNets/NNClassifier/gat_classifier.py b/photonai_graph/NeuralNets/NNClassifier/gat_classifier.py
deleted file mode 100644
index 77b3486..0000000
--- a/photonai_graph/NeuralNets/NNClassifier/gat_classifier.py
+++ /dev/null
@@ -1,71 +0,0 @@
-import numpy as np
-from typing import List
-from photonai_graph.NeuralNets.dgl_base import DGLmodel
-from photonai_graph.NeuralNets.NNModels import GATClassifier
-
-
-class GATClassifierModel(DGLmodel):
-
-    def __init__(self,
-                 in_dim: int = 1,
-                 hidden_layers: int = 2,
-                 hidden_dim: int = 256,
-                 heads: List= None,
-                 nn_epochs: int = 200,
-                 learning_rate: float = 0.001,
-                 batch_size: int = 32,
-                 adjacency_axis: int = 0,
-                 feature_axis: int = 1,
-                 logs: str = ''):
-        """
-            Graph Attention Network for graph classification. GAT Layers
-            are modeled after Veličković et al., 2018.
-            Implementation based on dgl & pytorch.
-
-
-            Parameters
-            ----------
-            in_dim: int,default=1
-                input dimension
-            hidden_layers: int,default=2
-                number of hidden layers used by the model
-            hidden_dim: int,default=256
-                dimensions in the hidden layers
-            heads: list,default=None
-                list with number of heads per hidden layer
-
-        """
-        super(GATClassifierModel, self).__init__(nn_epochs=nn_epochs,
-                                                 learning_rate=learning_rate,
-                                                 batch_size=batch_size,
-                                                 adjacency_axis=adjacency_axis,
-                                                 feature_axis=feature_axis,
-                                                 logs=logs)
-        if heads is None:
-            heads = [2, 2]
-            # Todo: if heads is not length of hidden layers +1 (bc of the first layer)
-        self.in_dim = in_dim
-        self.hidden_layers = hidden_layers
-        self.hidden_dim = hidden_dim
-        self.heads = heads
-
-    def fit(self, X, y):
-
-        # handle inputs
-        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
-        # get data loader
-        data_loader = self.get_data_loader(X_trans, y)
-        # specify model with optimizer etc
-        # set model class (import from NN.models)
-        self.model = GATClassifier(self.in_dim, self.hidden_dim, self.heads, len(np.unique(y)), self.hidden_layers)
-        # get optimizers
-        loss_func, optimizer = self.get_classifier()
-        # train model
-        self.model.train()
-        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
-
-        return self
-
-    def predict(self, x):
-
-        return self.predict_classifier(x)
diff --git a/photonai_graph/NeuralNets/NNClassifier/gcn_classifier.py b/photonai_graph/NeuralNets/NNClassifier/gcn_classifier.py
deleted file mode 100644
index f63bb63..0000000
--- a/photonai_graph/NeuralNets/NNClassifier/gcn_classifier.py
+++ /dev/null
@@ -1,74 +0,0 @@
-import numpy as np
-from photonai_graph.NeuralNets.dgl_base import DGLmodel
-from photonai_graph.NeuralNets.NNModels import GCNClassifier
-
-
-class GCNClassifierModel(DGLmodel):
-
-    def __init__(self,
-                 in_dim: int = 1,
-                 hidden_layers: int = 2,
-                 hidden_dim: int = 256,
-                 nn_epochs: int = 200,
-                 learning_rate: float = 0.001,
-                 batch_size: int = 32,
-                 adjacency_axis: int = 0,
-                 feature_axis: int = 1,
-                 allow_zero_in_degree: bool = False,
-                 logs: str = ''):
-        """
-        Graph Attention Network for graph classification. GCN Layers
-        from Kipf & Welling, 2017.
-        Implementation based on dgl & pytorch.
-
-
-        Parameters
-        ----------
-        in_dim: int,default=1
-            input dimension
-        hidden_layers: int,default=2
-            number of hidden layers used by the model
-        hidden_dim: int,default=256
-            dimensions in the hidden layers
-
-        """
-        super(GCNClassifierModel, self).__init__(nn_epochs=nn_epochs,
-                                                 learning_rate=learning_rate,
-                                                 batch_size=batch_size,
-                                                 adjacency_axis=adjacency_axis,
-                                                 feature_axis=feature_axis,
-                                                 allow_zero_in_degree=allow_zero_in_degree,
-                                                 logs=logs)
-        self.in_dim = in_dim
-        self.hidden_dim = hidden_dim
-        self.hidden_layers = hidden_layers
-
-    def fit(self, X, y):
-
-        # handle inputs
-        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
-        # get data loader
-        data_loader = self.get_data_loader(X_trans, y)
-        # specify model with optimizer etc
-        # set model class (import from NN.models)
-        self.model = GCNClassifier(self.in_dim,
-                                   self.hidden_dim,
-                                   len(np.unique(y)),
-                                   self.hidden_layers,
-                                   allow_zero_in_degree=self.allow_zero_in_degree)
-        # get optimizers
-        loss_func, optimizer = self.get_classifier()
-        # train model
-        self.model.train()
-        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
-
-        return self
-
-    def predict(self, X):
-
-        # self.model.eval()
-        # test_bg = dgl.batch(X)
-        # probs_y = torch.softmax(self.model(test_bg), 1)
-        # argmax_y = torch.max(probs_y, 1)[1].view(-1, 1)
-
-        return self.predict_classifier(X)
diff --git a/photonai_graph/NeuralNets/NNClassifier/sgc_classifier.py b/photonai_graph/NeuralNets/NNClassifier/sgc_classifier.py
deleted file mode 100644
index 3e2c996..0000000
--- a/photonai_graph/NeuralNets/NNClassifier/sgc_classifier.py
+++ /dev/null
@@ -1,62 +0,0 @@
-import numpy as np
-from photonai_graph.NeuralNets.dgl_base import DGLmodel
-from photonai_graph.NeuralNets.NNModels import SGConvClassifier
-
-
-class SGConvClassifierModel(DGLmodel):
-
-    def __init__(self,
-                 in_dim: int = 1,
-                 hidden_layers: int = 2,
-                 hidden_dim: int = 256,
-                 nn_epochs: int = 200,
-                 learning_rate: float = 0.001,
-                 batch_size: int = 32,
-                 adjacency_axis: int = 0,
-                 feature_axis: int = 1,
-                 logs: str = ''):
-        """
-        Graph convolutional network for graph classification. Simple Graph
-        convolutional layers from Wu, Felix, et al., 2018.
-        Implementation based on dgl & pytorch.
-
-
-        Parameters
-        ----------
-        in_dim: int,default=1
-            input dimension
-        hidden_layers: int,default=2
-            number of hidden layers used by the model
-        hidden_dim: int,default=256
-            dimensions in the hidden layers
-
-        """
-        super(SGConvClassifierModel, self).__init__(nn_epochs=nn_epochs,
-                                                    learning_rate=learning_rate,
-                                                    batch_size=batch_size,
-                                                    adjacency_axis=adjacency_axis,
-                                                    feature_axis=feature_axis,
-                                                    logs=logs)
-        self.in_dim = in_dim
-        self.hidden_layers = hidden_layers
-        self.hidden_dim = hidden_dim
-
-    def fit(self, X, y):
-
-        # handle inputs
-        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
-        # get data loader
-        data_loader = self.get_data_loader(X_trans, y)
-        # specify model with optimizer etc
-        self.model = SGConvClassifier(self.in_dim, self.hidden_dim, len(np.unique(y)), self.hidden_layers)
-        # get optimizers
-        loss_func, optimizer = self.get_classifier()
-        # train model
-        self.model.train()
-        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
-
-        return self
-
-    def predict(self, x):
-        """returns the argmax of the predictions"""
-        return self.predict_classifier(x)
diff --git a/photonai_graph/NeuralNets/NNLayers.py b/photonai_graph/NeuralNets/NNLayers.py
deleted file mode 100644
index f73ba50..0000000
--- a/photonai_graph/NeuralNets/NNLayers.py
+++ /dev/null
@@ -1,75 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import dgl.function as fn
-from dgl.nn.pytorch import edge_softmax
-
-
-class GATLayer(nn.Module):
-    def __init__(self,
-                 in_dim,
-                 out_dim,
-                 num_heads,
-                 feat_drop=0.,
-                 attn_drop=0.,
-                 alpha=0.2,
-                 agg_activation=F.elu):
-        super(GATLayer, self).__init__()
-
-        self.num_heads = num_heads
-        self.feat_drop = nn.Dropout(feat_drop)
-        self.fc = nn.Linear(in_dim, num_heads * out_dim, bias=False)
-        self.attn_l = nn.Parameter(torch.Tensor(size=(num_heads, out_dim, 1)))
-        self.attn_r = nn.Parameter(torch.Tensor(size=(num_heads, out_dim, 1)))
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.activation = nn.LeakyReLU(alpha)
-        self.softmax = edge_softmax
-
-        self.agg_activation = agg_activation
-
-    def clean_data(self):
-        ndata_names = ['ft', 'a1', 'a2']
-        edata_names = ['a_drop']
-        for name in ndata_names:
-            self.g.ndata.pop(name)
-        for name in edata_names:
-            self.g.edata.pop(name)
-
-    def forward(self, feat, bg):
-        # prepare, inputs are of shape V x F, V the number of nodes, F the dim of input features
-        self.g = bg
-        h = self.feat_drop(feat)
-        # V x K x F', K number of heads, F' dim of transformed features
-        ft = self.fc(h).reshape((h.shape[0], self.num_heads, -1))
-        head_ft = ft.transpose(0, 1)                              # K x V x F'
-        a1 = torch.bmm(head_ft, self.attn_l).transpose(0, 1)      # V x K x 1
-        a2 = torch.bmm(head_ft, self.attn_r).transpose(0, 1)      # V x K x 1
-        self.g.ndata.update({'ft': ft, 'a1': a1, 'a2': a2})
-        # 1. compute edge attention
-        self.g.apply_edges(self.edge_attention)
-        # 2. compute softmax in two parts: exp(x - max(x)) and sum(exp(x - max(x)))
-        self.edge_softmax()
-        # 2. compute the aggregated node features scaled by the dropped,
-        # unnormalized attention values.
-        self.g.update_all(fn.src_mul_edge('ft', 'a_drop', 'ft'), fn.sum('ft', 'ft'))
-        # 3. apply normalizer
-        ret = self.g.ndata['ft']                                  # V x K x F'
-        ret = ret.flatten(1)
-
-        if self.agg_activation is not None:
-            ret = self.agg_activation(ret)
-
-        # Clean ndata and edata
-        self.clean_data()
-
-        return ret
-
-    def edge_attention(self, edges):
-        # an edge UDF to compute un-normalized attention values from src and dst
-        a = self.activation(edges.src['a1'] + edges.dst['a2'])
-        return {'a': a}
-
-    def edge_softmax(self):
-        attention = self.softmax(self.g, self.g.edata.pop('a'))
-        # Dropout attention scores and save them
-        self.g.edata['a_drop'] = self.attn_drop(attention)
diff --git a/photonai_graph/NeuralNets/NNModels.py b/photonai_graph/NeuralNets/NNModels.py
deleted file mode 100644
index 90328c3..0000000
--- a/photonai_graph/NeuralNets/NNModels.py
+++ /dev/null
@@ -1,82 +0,0 @@
-import dgl
-import torch.nn as nn
-import torch.nn.functional as F
-from dgl.nn.pytorch import GraphConv, SGConv
-from photonai_graph.NeuralNets.NNLayers import GATLayer
-
-
-class GCNClassifier(nn.Module):
-    def __init__(self, in_dim, hidden_dim, n_classes, hidden_layers, allow_zero_in_degree):
-        super(GCNClassifier, self).__init__()
-        self.layers = nn.ModuleList()
-        # input layers
-        self.layers.append(GraphConv(in_dim, hidden_dim, allow_zero_in_degree=allow_zero_in_degree))
-        # hidden layers
-        for layer in range(1, hidden_layers):
-            self.layers.append(GraphConv(hidden_dim, hidden_dim, allow_zero_in_degree=allow_zero_in_degree))
-        self.classify = nn.Linear(hidden_dim, n_classes)
-
-    def forward(self, g):
-        # Use node degree as the initial node feature. For undirected graphs, the in-degree
-        # is the same as the out_degree.
-        h = g.in_degrees().view(-1, 1).float()
-        # Perform graph convolution and activation function.
-        for i, gnn in enumerate(self.layers):
-            h = F.relu(gnn(g, h))
-        g.ndata['h'] = h
-        # Calculate graph representation by averaging all the node representations.
-        hg = dgl.mean_nodes(g, 'h')
-        return self.classify(hg)
-
-
-class GATClassifier(nn.Module):
-    def __init__(self, in_dim, hidden_dim, num_heads, n_classes, hidden_layers):
-        super(GATClassifier, self).__init__()
-
-        self.layers = nn.ModuleList()
-        # append gat layer according to inputs
-        self.layers.append(GATLayer(in_dim, hidden_dim, num_heads[0]))
-        # hidden layers
-        for layer in range(1, hidden_layers):
-            self.layers.append(GATLayer(
-                hidden_dim * num_heads[layer - 1], hidden_dim, num_heads[layer]))
-        # output layer
-        self.classify = nn.Linear(hidden_dim * num_heads[-1], n_classes)
-
-    def forward(self, bg):
-        # For undirected graphs, in_degree is the same as
-        # out_degree.
-        h = bg.in_degrees().view(-1, 1).float()
-        for i, gnn in enumerate(self.layers):
-            h = gnn(h, bg)
-        bg.ndata['h'] = h
-        hg = dgl.mean_nodes(bg, 'h')
-        return self.classify(hg)
-
-
-class SGConvClassifier(nn.Module):
-    def __init__(self, in_dim, hidden_dim, n_classes, hidden_layers):
-        super(SGConvClassifier, self).__init__()
-        self.in_dim = in_dim
-        self.hidden_dim = hidden_dim
-        self.n_classes = n_classes
-        self.hidden_layers = hidden_layers
-        self.layers = nn.ModuleList()
-
-        # input layer
-        self.layers.append(SGConv(in_dim, hidden_dim))
-        # hidden layers
-        for lr in range(1, hidden_layers):
-            self.layers.append(SGConv(hidden_dim, hidden_dim))
-        # output layer
-        self.classify = nn.Linear(hidden_dim, n_classes)
-
-    def forward(self, bg):
-        h = bg.in_degrees().view(-1, 1).float()
-        for lr, layer in enumerate(self.layers):
-            h = layer(bg, h)
-        bg.ndata['h'] = h
-        # Calculate graph representation by averaging all the node representations.
-        hg = dgl.mean_nodes(bg, 'h')
-
-        return self.classify(hg)
diff --git a/photonai_graph/NeuralNets/NNRegressor/__init__.py b/photonai_graph/NeuralNets/NNRegressor/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/photonai_graph/NeuralNets/NNRegressor/gat_regressor.py b/photonai_graph/NeuralNets/NNRegressor/gat_regressor.py
deleted file mode 100644
index 376cc21..0000000
--- a/photonai_graph/NeuralNets/NNRegressor/gat_regressor.py
+++ /dev/null
@@ -1,70 +0,0 @@
-from typing import List
-
-from photonai_graph.NeuralNets.dgl_base import DGLmodel
-from photonai_graph.NeuralNets.NNModels import GATClassifier
-
-
-class GATRegressorModel(DGLmodel):
-
-    def __init__(self,
-                 in_dim: int = 1,
-                 hidden_layers: int = 2,
-                 hidden_dim: int = 256,
-                 heads: List  = None,
-                 nn_epochs: int = 200,
-                 learning_rate: float = 0.001,
-                 batch_size: int = 32,
-                 adjacency_axis: int = 0,
-                 feature_axis: int = 1,
-                 logs: str = ''):
-        """
-            Graph Attention Network for graph regression. GAT Layers
-            are modeled after Veličković et al., 2018. Implementation
-            based on dgl & pytorch.
-
-
-            Parameters
-            ----------
-            in_dim: int,default=1
-                input dimension
-            hidden_layers: int,default=2
-                number of hidden layers used by the model
-            hidden_dim: int,default=256
-                dimensions in the hidden layers
-            heads: list,default=None
-                list with number of heads per hidden layer
-
-        """
-        super(GATRegressorModel, self).__init__(nn_epochs=nn_epochs,
-                                                learning_rate=learning_rate,
-                                                batch_size=batch_size,
-                                                adjacency_axis=adjacency_axis,
-                                                feature_axis=feature_axis,
-                                                logs=logs)
-        if heads is None:
-            heads = [2, 2]
-            # Todo: if heads is not length of hidden layers +1 (bc of the first layer)
-        self.in_dim = in_dim
-        self.hidden_layers = hidden_layers
-        self.hidden_dim = hidden_dim
-        self.heads = heads
-
-    def fit(self, X, y):
-
-        # handle inputs
-        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
-        # get data loader
-        data_loader = self.get_data_loader_regression(X_trans, y)
-        # specify model with optimizer etc
-        self.model = GATClassifier(self.in_dim, self.hidden_dim, self.heads, 1, self.hidden_layers).float()
-        # get optimizers
-        loss_func, optimizer = self.get_regressor()
-        # train model
-        self.model.train()
-        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
-
-        return self
-
-    def predict(self, x):
-
-        return self.predict_regressor(x)
diff --git a/photonai_graph/NeuralNets/NNRegressor/gcn_regressor.py b/photonai_graph/NeuralNets/NNRegressor/gcn_regressor.py
deleted file mode 100644
index fa6111d..0000000
--- a/photonai_graph/NeuralNets/NNRegressor/gcn_regressor.py
+++ /dev/null
@@ -1,65 +0,0 @@
-from photonai_graph.NeuralNets.dgl_base import DGLmodel
-from photonai_graph.NeuralNets.NNModels import GCNClassifier
-
-
-class GCNRegressorModel(DGLmodel):
-
-    def __init__(self,
-                 in_dim: int = 1,
-                 hidden_layers: int = 2,
-                 hidden_dim: int = 256,
-                 nn_epochs: int = 200,
-                 learning_rate: float = 0.001,
-                 batch_size: int = 32,
-                 adjacency_axis: int = 0,
-                 allow_zero_in_degree: bool = False,
-                 feature_axis: int = 1,
-                 logs: str = ''):
-        """
-        Graph convolutional Network for graph regression. GCN Layers
-        from Kipf & Welling, 2017. Implementation based on dgl & pytorch.
-
-
-        Parameters
-        ----------
-        in_dim: int,default=1
-            input dimension
-        hidden_layers: int,default=2
-            number of hidden layers used by the model
-        hidden_dim: int,default=256
-            dimensions in the hidden layers
-
-        """
-        super(GCNRegressorModel, self).__init__(nn_epochs=nn_epochs,
-                                                learning_rate=learning_rate,
-                                                batch_size=batch_size,
-                                                adjacency_axis=adjacency_axis,
-                                                feature_axis=feature_axis,
-                                                allow_zero_in_degree=allow_zero_in_degree,
-                                                logs=logs)
-        self.in_dim = in_dim
-        self.hidden_dim = hidden_dim
-        self.hidden_layers = hidden_layers
-
-    def fit(self, X, y):
-
-        # handle inputs
-        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
-        # get data loader
-        data_loader = self.get_data_loader_regression(X_trans, y)
-        # specify model with optimizer etc
-        self.model = GCNClassifier(self.in_dim,
-                                   self.hidden_dim, 1,
-                                   self.hidden_layers,
-                                   allow_zero_in_degree=self.allow_zero_in_degree).float()
-        # get optimizers
-        loss_func, optimizer = self.get_regressor()
-        # train model
-        self.model.train()
-        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
-
-        return self
-
-    def predict(self, X):
-
-        return self.predict_regressor(X)
diff --git a/photonai_graph/NeuralNets/NNRegressor/sgc_regressor.py b/photonai_graph/NeuralNets/NNRegressor/sgc_regressor.py
deleted file mode 100644
index 8e93160..0000000
--- a/photonai_graph/NeuralNets/NNRegressor/sgc_regressor.py
+++ /dev/null
@@ -1,61 +0,0 @@
-from photonai_graph.NeuralNets.dgl_base import DGLmodel
-from photonai_graph.NeuralNets.NNModels import SGConvClassifier
-
-
-class SGConvRegressorModel(DGLmodel):
-
-    def __init__(self,
-                 in_dim: int = 1,
-                 hidden_layers: int = 2,
-                 hidden_dim: int = 256,
-                 nn_epochs: int = 200,
-                 learning_rate: float = 0.001,
-                 batch_size: int = 32,
-                 adjacency_axis: int = 0,
-                 feature_axis: int = 1,
-                 logs: str = ''):
-        """
-        Graph convolutional network for graph regression. Simple Graph
-        convolutional layers from Wu, Felix, et al., 2018. Implementation
-        based on dgl & pytorch.
-
-
-        Parameters
-        ----------
-        in_dim: int,default=1
-            input dimension
-        hidden_layers: int,default=2
-            number of hidden layers used by the model
-        hidden_dim: int,default=256
-            dimensions in the hidden layers
-
-        """
-        super(SGConvRegressorModel, self).__init__(nn_epochs=nn_epochs,
-                                                   learning_rate=learning_rate,
-                                                   batch_size=batch_size,
-                                                   adjacency_axis=adjacency_axis,
-                                                   feature_axis=feature_axis,
-                                                   logs=logs)
-        self.in_dim = in_dim
-        self.hidden_layers = hidden_layers
-        self.hidden_dim = hidden_dim
-
-    def fit(self, X, y):
-
-        # handle inputs
-        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
-        # get data loader
-        data_loader = self.get_data_loader_regression(X_trans, y)
-        # specify model with optimizer etc
-        self.model = SGConvClassifier(self.in_dim, self.hidden_dim, 1, self.hidden_layers).float()
-        # get optimizers
-        loss_func, optimizer = self.get_regressor()
-        # train model
-        self.model.train()
-        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
-
-        return self
-
-    def predict(self, x):
-        """returns the argmax of the predictions"""
-        return self.predict_regressor(x)
diff --git a/photonai_graph/NeuralNets/SGCModel.py b/photonai_graph/NeuralNets/SGCModel.py
new file mode 100644
index 0000000..d10a50a
--- /dev/null
+++ b/photonai_graph/NeuralNets/SGCModel.py
@@ -0,0 +1,129 @@
+import dgl
+import numpy as np
+import torch.nn as nn
+from dgl.nn.pytorch import GraphConv, SGConv, GATConv
+
+from photonai_graph.NeuralNets.dgl_base import DGLClassifierBaseModel, DGLRegressorBaseModel
+
+
+class SGConvClassifier(nn.Module):
+    def __init__(self, in_dim, hidden_dim, n_classes, hidden_layers, allow_zero_in_degree):
+        super(SGConvClassifier, self).__init__()
+        self.in_dim = in_dim
+        self.hidden_dim = hidden_dim
+        self.n_classes = n_classes
+        self.hidden_layers = hidden_layers
+        self.layers = nn.ModuleList()
+
+        # input layer
+        self.layers.append(SGConv(in_dim, hidden_dim))
+        # hidden layers
+        for lr in range(1, hidden_layers):
+            self.layers.append(SGConv(hidden_dim, hidden_dim, allow_zero_in_degree=allow_zero_in_degree))
+        # output layer
+        self.classify = nn.Linear(hidden_dim, n_classes)
+
+    def forward(self, bg):
+        h = bg.in_degrees().view(-1, 1).float()
+        for lr, layer in enumerate(self.layers):
+            h = layer(bg, h)
+        bg.ndata['h'] = h
+        # Calculate graph representation by averaging all the node representations.
+        hg = dgl.mean_nodes(bg, 'h')
+
+        return self.classify(hg)
+
+
+class SGConvClassifierModel(DGLClassifierBaseModel):
+
+    def __init__(self,
+                 in_dim: int = 1,
+                 hidden_layers: int = 2,
+                 hidden_dim: int = 256,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = ''):
+        """
+        Graph convolutional network for graph classification. Simple Graph
+        convolutional layers from Wu, Felix, et al., 2018.
+        Implementation based on dgl & pytorch.
+
+
+        Parameters
+        ----------
+        in_dim: int,default=1
+            input dimension
+        hidden_layers: int,default=2
+            number of hidden layers used by the model
+        hidden_dim: int,default=256
+            dimensions in the hidden layers
+
+        """
+        super(SGConvClassifierModel, self).__init__(nn_epochs=nn_epochs,
+                                                    learning_rate=learning_rate,
+                                                    batch_size=batch_size,
+                                                    adjacency_axis=adjacency_axis,
+                                                    feature_axis=feature_axis,
+                                                    add_self_loops=add_self_loops,
+                                                    allow_zero_in_degree=allow_zero_in_degree,
+                                                    logs=logs)
+        self.in_dim = in_dim
+        self.hidden_layers = hidden_layers
+        self.hidden_dim = hidden_dim
+
+    def _init_model(self, X=None, y=None):
+        self.model = SGConvClassifier(self.in_dim, self.hidden_dim,
+                                      len(np.unique(y)), self.hidden_layers,
+                                      allow_zero_in_degree=self.allow_zero_in_degree)
+
+
+class SGConvRegressorModel(DGLRegressorBaseModel):
+
+    def __init__(self,
+                 in_dim: int = 1,
+                 hidden_layers: int = 2,
+                 hidden_dim: int = 256,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = ''):
+        """
+        Graph convolutional network for graph regression. Simple Graph
+        convolutional layers from Wu, Felix, et al., 2018. Implementation
+        based on dgl & pytorch.
+
+
+        Parameters
+        ----------
+        in_dim: int,default=1
+            input dimension
+        hidden_layers: int,default=2
+            number of hidden layers used by the model
+        hidden_dim: int,default=256
+            dimensions in the hidden layers
+
+        """
+        super(SGConvRegressorModel, self).__init__(nn_epochs=nn_epochs,
+                                                   learning_rate=learning_rate,
+                                                   batch_size=batch_size,
+                                                   adjacency_axis=adjacency_axis,
+                                                   feature_axis=feature_axis,
+                                                   add_self_loops=add_self_loops,
+                                                   allow_zero_in_degree=allow_zero_in_degree,
+                                                   logs=logs)
+        self.in_dim = in_dim
+        self.hidden_layers = hidden_layers
+        self.hidden_dim = hidden_dim
+
+    def _init_model(self, X=None, y=None):
+        self.model = SGConvClassifier(self.in_dim, self.hidden_dim, 1, self.hidden_layers,
+                                      allow_zero_in_degree=self.allow_zero_in_degree).float()
diff --git a/photonai_graph/NeuralNets/dgl_base.py b/photonai_graph/NeuralNets/dgl_base.py
index 1ef0685..b514d57 100644
--- a/photonai_graph/NeuralNets/dgl_base.py
+++ b/photonai_graph/NeuralNets/dgl_base.py
@@ -1,8 +1,9 @@
 import warnings
-from abc import ABC
+from abc import ABC, abstractmethod
 import os
 from photonai_graph.util import assert_imported
-from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
+from tqdm import tqdm
 try:
     import dgl
     import torch
@@ -16,7 +17,7 @@
     pass
 
 
-class DGLmodel(BaseEstimator, ClassifierMixin, ABC):
+class DGLModel(BaseEstimator, ABC):
     _estimator_type = "classifier"
 
     def __init__(self, nn_epochs: int = 200,
@@ -26,11 +27,12 @@ def __init__(self, nn_epochs: int = 200,
                  feature_axis: int = 1,
                  add_self_loops: bool = True,
                  allow_zero_in_degree: bool = False,
+                 verbose: bool = False,
                  logs: str = None):
         """
         Base class for DGL based graph neural networks. Implements
         helper functions and shared parameters used by other models.
-        Implementation based on gem python package.
+        Implementation based on dgl python package.
 
 
         Parameters
@@ -49,6 +51,8 @@ def __init__(self, nn_epochs: int = 200,
             self loops are added if true
         allow_zero_in_degree: bool,default=False
             If true the zero in degree test of dgl is disabled
+        verbose: bool,default=False
+            If true verbose information is printed
         logs: str,default=None
             Path to the log data
 
@@ -60,20 +64,23 @@ def __init__(self, nn_epochs: int = 200,
         self.feature_axis = feature_axis
         self.add_self_loops = add_self_loops
         self.allow_zero_in_degree = allow_zero_in_degree
-        if self.add_self_loops and not self.allow_zero_in_degree:
+        if self.add_self_loops and self.allow_zero_in_degree:
             warnings.warn('If self loops are added allow_zero_in_degree should be false!')
+        if not self.add_self_loops and not self.allow_zero_in_degree:
+            warnings.warn('If no self loops are added and allow_zero_in_degree is set to false, '
+                          'all graphs should not contain 0-in-degree nodes')
         self.model = None
         if logs:
             self.logs = logs
         else:
             self.logs = os.getcwd()
         assert_imported(["dgl", "pytorch"])
+        self.verbose = verbose
 
-    @staticmethod
-    def train_model(epochs, model, optimizer, loss_func, data_loader):
+    def train_model(self, epochs, model, optimizer, loss_func, data_loader):
         # This function trains the neural network
         epoch_losses = []
-        for epoch in range(epochs):
+        for epoch in tqdm(range(epochs)):
             epoch_loss = 0
             iteration = 0
             for it, (bg, label) in enumerate(data_loader):
@@ -85,78 +92,190 @@ def train_model(epochs, model, optimizer, loss_func, data_loader):
                 epoch_loss += loss.detach().item()
                 iteration = it
             epoch_loss /= (iteration + 1)
-            print('Epoch {}, loss {:.4f}'.format(epoch, epoch_loss))
+            if self.verbose:
+                print('Epoch {}, loss {:.4f}'.format(epoch, epoch_loss))
             epoch_losses.append(epoch_loss)
 
-    @staticmethod
-    def collate(samples):
-        """returns a batched graph, the input (samples) is a list of pairs (graph, label)"""
-        graphs, labels = map(list, zip(*samples))
-        batched_graph = dgl.batch(graphs)
-        return batched_graph, torch.tensor(labels, dtype=torch.long)
-
-    @staticmethod
-    def collate_regression(samples):
-        """returns a batched graph, the input (samples) is a list of pairs"""
-        graphs, labels = map(list, zip(*samples))
-        batched_graph = dgl.batch(graphs)
-        return batched_graph, torch.tensor(labels, dtype=torch.float32)
-
-    @staticmethod
-    def handle_inputs(x, adjacency_axis, feature_axis):
+    def handle_inputs(self, x, adjacency_axis, feature_axis):
         """checks the format of the input and transforms them for dgl models"""
         x_trans = check_dgl(x, adjacency_axis=adjacency_axis, feature_axis=feature_axis)
+        if self.add_self_loops:
+            x_trans = [dgl.add_self_loop(x) for x in x_trans]
         return x_trans
 
+    def fit(self, X, y=None):
+        # handle inputs
+        X_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
+        # get data loader
+        data_loader = self.get_data_loader(X_trans, y)
+        # specify model with optimizer etc
+        self._init_model(X, y)
+        # get optimizers
+        loss_func, optimizer = self.setup_model()
+        # train model
+        self.model.train()
+        self.train_model(self.nn_epochs, self.model, optimizer, loss_func, data_loader)
+        return self
+
+    def predict(self, x):
+        return self.predict_model(x)
+
+    @staticmethod
+    @abstractmethod
+    def collate(samples):
+        """Collate function"""
+
+    @abstractmethod
     def get_data_loader(self, x_trans, y):
-        """returns data in a data loader format"""
-        # add a self loop for 0-in-degree nodes
-        if self.add_self_loops:
-            x_san = [dgl.add_self_loop(x) for x in x_trans]
-        # create dataloader
-        data = DGLData(zip_data(x_san, y))
-        data_loader = GraphDataLoader(data, batch_size=self.batch_size, shuffle=True, collate_fn=self.collate)
+        """Data loader"""
 
-        return data_loader
+    @abstractmethod
+    def predict_model(self, X):
+        """get model predictions"""
 
-    def get_data_loader_regression(self, x_trans, y):
-        """returns data in a regression data loader format"""
-        y = y.reshape(y.shape[0], 1)
-        data = DGLData(zip_data(x_trans, y))
-        data_loader = DataLoader(data, batch_size=self.batch_size, shuffle=True, collate_fn=self.collate_regression)
+    @abstractmethod
+    def setup_model(self):
+        """Setup the model"""
 
-        return data_loader
+    @abstractmethod
+    def _init_model(self, X=None, y=None):
+        """initialize model"""
+
+
+class DGLClassifierBaseModel(DGLModel, ClassifierMixin, ABC):
+    def __init__(self,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = None):
+        """Abstract base class for classification algorithms
 
-    def predict_classifier(self, x):
+                Parameters
+                ----------
+                nn_epochs: int,default=200
+                    Number of epochs to fit the model
+                learning_rate: float,default=0.001
+                    Learning rate for model training
+                batch_size: int,default=32
+                    Batch size for model training
+                adjacency_axis: int,default=0
+                    Axis which contains the adjacency
+                feature_axis: int,default=1
+                    Axis which contains the features
+                add_self_loops: bool,default=True
+                    If this value is true, a self loop is added to each node of each graph
+                allow_zero_in_degree: bool,default=False
+                    If true the dgl model allows zero-in-degree Graphs
+                logs: str,default=None
+                    Default logging directory
+                """
+        super(DGLClassifierBaseModel, self).__init__(nn_epochs=nn_epochs,
+                                                     learning_rate=learning_rate,
+                                                     batch_size=batch_size,
+                                                     adjacency_axis=adjacency_axis,
+                                                     feature_axis=feature_axis,
+                                                     add_self_loops=add_self_loops,
+                                                     allow_zero_in_degree=allow_zero_in_degree,
+                                                     logs=logs)
+
+    def setup_model(self):
+        """returns the loss and optimizer for classification"""
+        loss_func = nn.CrossEntropyLoss()  # specify loss
+        optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
+        return loss_func, optimizer
+
+    def predict_model(self, X):
         """returns the argmax of the predictions for classification tasks"""
         self.model.eval()
-        x_trans = self.handle_inputs(x, self.adjacency_axis, self.feature_axis)
+        x_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
         test_bg = dgl.batch(x_trans)
         probs_y = torch.softmax(self.model(test_bg), 1)
         argmax_y = torch.max(probs_y, 1)[1].view(-1, 1)
+        return argmax_y.squeeze()
 
-        return argmax_y
+    def get_data_loader(self, x_trans, y):
+        """returns data in a data loader format"""
+        data = DGLData(zip_data(x_trans, y))
+        # create dataloader
+        data_loader = GraphDataLoader(data, batch_size=self.batch_size, shuffle=True, collate_fn=self.collate)
+        return data_loader
 
-    def predict_regressor(self, x):
-        """returns the predictions for a regression model"""
-        self.model.eval()
-        x_trans = self.handle_inputs(x, self.adjacency_axis, self.feature_axis)
-        test_bg = dgl.batch(x_trans)
-        probs = self.model(test_bg)
-        probs = probs.detach().numpy()
+    @staticmethod
+    def collate(samples):
+        """returns a batched graph, the input (samples) is a list of pairs (graph, label)"""
+        graphs, labels = map(list, zip(*samples))
+        batched_graph = dgl.batch(graphs)
+        return batched_graph, torch.tensor(labels, dtype=torch.long)
 
-        return probs
 
-    def get_classifier(self):
-        """returns the loss and optimizer for classification"""
-        loss_func = nn.CrossEntropyLoss()  # specify loss
-        optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
+class DGLRegressorBaseModel(DGLModel, RegressorMixin, ABC):
+    def __init__(self,
+                 nn_epochs: int = 200,
+                 learning_rate: float = 0.001,
+                 batch_size: int = 32,
+                 adjacency_axis: int = 0,
+                 feature_axis: int = 1,
+                 add_self_loops: bool = True,
+                 allow_zero_in_degree: bool = False,
+                 logs: str = None):
+        """Abstract base class for regression algorithms
 
-        return loss_func, optimizer
+        Parameters
+        ----------
+        nn_epochs: int,default=200
+            Number of epochs to fit the model
+        learning_rate: float,default=0.001
+            Learning rate for model training
+        batch_size: int,default=32
+            Batch size for model training
+        adjacency_axis: int,default=0
+            Axis which contains the adjacency
+        feature_axis: int,default=1
+            Axis which contains the features
+        add_self_loops: bool,default=True
+            If this value is true, a self loop is added to each node of each graph
+        allow_zero_in_degree: bool,default=False
+            If true the dgl model allows zero-in-degree Graphs
+        logs: str,default=None
+            Default logging directory
+        """
+        super(DGLRegressorBaseModel, self).__init__(nn_epochs=nn_epochs,
+                                                    learning_rate=learning_rate,
+                                                    batch_size=batch_size,
+                                                    adjacency_axis=adjacency_axis,
+                                                    feature_axis=feature_axis,
+                                                    add_self_loops=add_self_loops,
+                                                    allow_zero_in_degree=allow_zero_in_degree,
+                                                    logs=logs)
 
-    def get_regressor(self):
+    def setup_model(self):
         """returns the loss and optimizer for regression"""
         loss_func = nn.MSELoss()  # regression loss
         optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
-
         return loss_func, optimizer
+
+    def predict_model(self, X):
+        self.model.eval()
+        x_trans = self.handle_inputs(X, self.adjacency_axis, self.feature_axis)
+        test_bg = dgl.batch(x_trans)
+        probs = self.model(test_bg)
+        probs = probs.detach().numpy()
+        return probs.squeeze()
+
+    def get_data_loader(self, x_trans, y):
+        """returns data in a regression data loader format"""
+        y = y.reshape(y.shape[0], 1)
+        data = DGLData(zip_data(x_trans, y))
+        data_loader = DataLoader(data, batch_size=self.batch_size, shuffle=True, collate_fn=self.collate)
+        return data_loader
+
+    @staticmethod
+    def collate(samples):
+        """returns a batched graph, the input (samples) is a list of pairs"""
+        graphs, labels = map(list, zip(*samples))
+        batched_graph = dgl.batch(graphs)
+        return batched_graph, torch.tensor(labels, dtype=torch.float32)
diff --git a/photonai_graph/photonai_graph.json b/photonai_graph/photonai_graph.json
index fcec879..e3a96cd 100644
--- a/photonai_graph/photonai_graph.json
+++ b/photonai_graph/photonai_graph.json
@@ -56,27 +56,27 @@
       "Transformer"
    ],
    "GATClassifier":[
-      "photonai_graph.NeuralNets.NNClassifier.gat_classifier.GATClassifierModel",
+      "photonai_graph.NeuralNets.GATModel.GATClassifierModel",
       "Estimator"
    ],
    "GCNClassifier":[
-      "photonai_graph.NeuralNets.NNClassifier.gcn_classifier.GCNClassifierModel",
+      "photonai_graph.NeuralNets.GCNModel.GCNClassifierModel",
       "Estimator"
    ],
    "SGCClassifier":[
-      "photonai_graph.NeuralNets.NNClassifier.sgc_classifier.SGConvClassifierModel",
+      "photonai_graph.NeuralNets.SGCModel.SGConvClassifierModel",
       "Estimator"
    ],
    "GATRegressor": [
-      "photonai_graph.NeuralNets.NNRegressor.gat_regressor.GATRegressorModel",
+      "photonai_graph.NeuralNets.GATModel.GATRegressorModel",
       "Estimator"
    ],
    "GCNRegressor": [
-      "photonai_graph.NeuralNets.NNRegressor.gcn_regressor.GCNRegressorModel",
+      "photonai_graph.NeuralNets.GCNModel.GCNRegressorModel",
       "Estimator"
    ],
    "SGCRegressor": [
-      "photonai_graph.NeuralNets.NNRegressor.sgc_regressor.SGCRegressorModel",
+      "photonai_graph.NeuralNets.SGCModel.SGCRegressorModel",
       "Estimator"
    ],
    "GrakelAdapter": [
diff --git a/requirements.txt b/requirements.txt
index f3d2a01..7342730 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,4 +9,5 @@ pandas
 numpy
 scipy
 tensorflow
-torch
\ No newline at end of file
+torch
+tqdm
\ No newline at end of file
diff --git a/setup.py b/setup.py
index f13f9fc..78f0a0e 100644
--- a/setup.py
+++ b/setup.py
@@ -29,5 +29,6 @@
                       'nxt_gem @ git+https://github.com/jernsting/nxt_gem.git',
                       'pandas>=1.1.2',
                       'numpy>=1.12.2',
-                      'scipy>=0.19.1']
+                      'scipy>=0.19.1',
+                      'tqdm']
 )


Test patch:
diff --git a/.github/workflows/lintandtest.yml b/.github/workflows/lintandtest.yml
index 3a3675d..4202d81 100644
--- a/.github/workflows/lintandtest.yml
+++ b/.github/workflows/lintandtest.yml
@@ -5,9 +5,9 @@ name: Python application
 
 on:
   push:
-    branches: [ master ]
+    branches: [ master, dev ]
   pull_request:
-    branches: [ master ]
+    branches: [ master, dev ]
 
 jobs:
   build:
diff --git a/test/conversion_tests/test_dense_to_dgl.py b/test/conversion_tests/test_dense_to_dgl.py
index 2f80ad6..21fbfc1 100644
--- a/test/conversion_tests/test_dense_to_dgl.py
+++ b/test/conversion_tests/test_dense_to_dgl.py
@@ -21,3 +21,14 @@ def test_matrix(self):
     def test_nonsense_input(self):
         with self.assertRaises(ValueError):
             dense_to_dgl(self.edge_dict, adjacency_axis=0, feature_axis=1)
+
+    def test_disconnected_graph(self):
+        in_graph = np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0]])
+        in_feats = np.array([[.9, .3, .2], [.8, .2, .1], [.7, .1, .0]])
+        in_graph = np.expand_dims(in_graph, axis=(0, -1))
+        in_feats = np.expand_dims(in_feats, axis=(0, -1))
+        in_graph = np.concatenate((in_graph, in_feats), axis=-1)
+        dgl_graph = dense_to_dgl(in_graph, adjacency_axis=0, feature_axis=1)[0]
+        self.assertEqual(dgl_graph.num_nodes(), 3)
+        self.assertEqual(dgl_graph.num_edges(), 2)
+        self.assertTrue(np.array_equal(dgl_graph.ndata['feat'].numpy(), in_feats[0, ..., 0]))
diff --git a/photonai_graph/NeuralNets/NNClassifier/__init__.py b/test/io_tests/__init__.py
similarity index 100%
rename from photonai_graph/NeuralNets/NNClassifier/__init__.py
rename to test/io_tests/__init__.py
diff --git a/test/conversion_tests/test_load_file_to_networkx.py b/test/io_tests/test_load_file_to_networkx.py
similarity index 100%
rename from test/conversion_tests/test_load_file_to_networkx.py
rename to test/io_tests/test_load_file_to_networkx.py
diff --git a/test/conversion_tests/test_save_networkx_to_file.py b/test/io_tests/test_save_networkx_to_file.py
similarity index 96%
rename from test/conversion_tests/test_save_networkx_to_file.py
rename to test/io_tests/test_save_networkx_to_file.py
index fe99a88..b812dd9 100644
--- a/test/conversion_tests/test_save_networkx_to_file.py
+++ b/test/io_tests/test_save_networkx_to_file.py
@@ -2,7 +2,8 @@
 import numpy as np
 import networkx as nx
 import os
-from photonai_graph.GraphConversions import save_networkx_to_file
+
+from photonai_graph.GraphIO import save_networkx_to_file
 
 
 class SaveNetworkxToFileTest(unittest.TestCase):
diff --git a/test/nn_tests/test_gat_classifier.py b/test/nn_tests/test_gat_classifier.py
index 0d9f49e..c2cdcba 100644
--- a/test/nn_tests/test_gat_classifier.py
+++ b/test/nn_tests/test_gat_classifier.py
@@ -2,7 +2,7 @@
 import numpy as np
 import networkx as nx
 from dgl.data import MiniGCDataset
-from photonai_graph.NeuralNets.NNClassifier.gat_classifier import GATClassifierModel
+from photonai_graph.NeuralNets.GATModel import GATClassifierModel
 from photonai_graph.GraphUtilities import get_random_labels
 
 
@@ -25,22 +25,22 @@ def test_gat_classifier_output_shape(self):
         gat_clf = GATClassifierModel(nn_epochs=20)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_gat_classifier_output_hidden_dim(self):
         gat_clf = GATClassifierModel(hidden_dim=128)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_gat_classifier_nx_graphs(self):
         gat_clf = GATClassifierModel(nn_epochs=20)
         gat_clf.fit(self.X_nx, self.y)
         output = gat_clf.predict(self.X_nx)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_gat_classifier_dgl(self):
         gat_clf = GATClassifierModel(nn_epochs=20)
         gat_clf.fit(self.X_dgl, self.y)
         output = gat_clf.predict(self.X_dgl)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
diff --git a/test/nn_tests/test_gat_regressor.py b/test/nn_tests/test_gat_regressor.py
index 9833bc2..898a016 100644
--- a/test/nn_tests/test_gat_regressor.py
+++ b/test/nn_tests/test_gat_regressor.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 import networkx as nx
-from photonai_graph.NeuralNets.NNRegressor.gat_regressor import GATRegressorModel
+from photonai_graph.NeuralNets.GATModel import GATRegressorModel
 from photonai_graph.GraphUtilities import get_random_labels
 
 
@@ -9,7 +9,7 @@ class GATRegressorTests(unittest.TestCase):
 
     def setUp(self):
         mtrx = np.random.rand(20, 20, 20, 2)
-        mtrx[:, :, :, 0][mtrx[:, :, :, 0] <= 0.9] = 0
+        # mtrx[:, :, :, 0][mtrx[:, :, :, 0] <= 0.9] = 0
         self.Xrandom4d = mtrx
         self.X_nx = [nx.erdos_renyi_graph(20, p=0.3)] * 20
         self.y = get_random_labels(number_of_labels=20)
@@ -18,10 +18,10 @@ def test_gat_regressor_output_shape(self):
         gat_clf = GATRegressorModel(nn_epochs=20)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_gat_regressor_output_hidden_dim(self):
         gat_clf = GATRegressorModel(hidden_dim=128)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
diff --git a/test/nn_tests/test_gcn_classifier.py b/test/nn_tests/test_gcn_classifier.py
index fd0de9b..9bba8ab 100644
--- a/test/nn_tests/test_gcn_classifier.py
+++ b/test/nn_tests/test_gcn_classifier.py
@@ -2,7 +2,7 @@
 import numpy as np
 import networkx as nx
 from dgl.data import MiniGCDataset
-from photonai_graph.NeuralNets.NNClassifier.gcn_classifier import GCNClassifierModel
+from photonai_graph.NeuralNets.GCNModel import GCNClassifierModel
 from photonai_graph.GraphUtilities import get_random_labels
 
 
@@ -25,22 +25,22 @@ def test_gcn_classifier_output_shape(self):
         gcn_clf = GCNClassifierModel(nn_epochs=20)
         gcn_clf.fit(self.Xrandom4d, self.y)
         output = gcn_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertTrue(np.array_equal(output.shape, self.y.shape))
 
     def test_gcn_classifier_output_hidden_dim(self):
         gcn_clf = GCNClassifierModel(hidden_dim=128)
         gcn_clf.fit(self.Xrandom4d, self.y)
         output = gcn_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertTrue(np.array_equal(output.shape, self.y.shape))
 
     def test_gat_classifier_nx_graphs(self):
         gat_clf = GCNClassifierModel(nn_epochs=20)
         gat_clf.fit(self.X_nx, self.y)
         output = gat_clf.predict(self.X_nx)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertTrue(np.array_equal(output.shape, self.y.shape))
 
     def test_gat_classifier_dgl(self):
         gat_clf = GCNClassifierModel(nn_epochs=20)
         gat_clf.fit(self.X_dgl, self.y)
         output = gat_clf.predict(self.X_dgl)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertTrue(np.array_equal(np.array(output.shape), self.y.shape))
diff --git a/test/nn_tests/test_gcn_regressor.py b/test/nn_tests/test_gcn_regressor.py
index fe2de09..2d48377 100644
--- a/test/nn_tests/test_gcn_regressor.py
+++ b/test/nn_tests/test_gcn_regressor.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 import networkx as nx
-from photonai_graph.NeuralNets.NNRegressor.gcn_regressor import GCNRegressorModel
+from photonai_graph.NeuralNets.GCNModel import GCNRegressorModel
 from photonai_graph.GraphUtilities import get_random_labels
 
 
@@ -20,10 +20,10 @@ def test_gcn_regressor_output_shape(self):
         gat_clf = GCNRegressorModel(nn_epochs=20)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_gcn_regressor_output_hidden_dim(self):
         gat_clf = GCNRegressorModel(hidden_dim=128)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
diff --git a/test/nn_tests/test_sgc_classifier.py b/test/nn_tests/test_sgc_classifier.py
index fbbfa24..61542a9 100644
--- a/test/nn_tests/test_sgc_classifier.py
+++ b/test/nn_tests/test_sgc_classifier.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 import networkx as nx
-from photonai_graph.NeuralNets.NNClassifier.sgc_classifier import SGConvClassifierModel
+from photonai_graph.NeuralNets.SGCModel import SGConvClassifierModel
 from photonai_graph.GraphUtilities import get_random_labels
 
 
@@ -21,10 +21,10 @@ def test_sgc_classifier_output_shape(self):
         gat_clf = SGConvClassifierModel(nn_epochs=20)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_sgc_classifier_output_hidden_dim(self):
         gat_clf = SGConvClassifierModel(hidden_dim=128)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
diff --git a/test/nn_tests/test_sgc_regressor.py b/test/nn_tests/test_sgc_regressor.py
index 7b7f14a..dab7c0a 100644
--- a/test/nn_tests/test_sgc_regressor.py
+++ b/test/nn_tests/test_sgc_regressor.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 import networkx as nx
-from photonai_graph.NeuralNets.NNRegressor.sgc_regressor import SGConvRegressorModel
+from photonai_graph.NeuralNets.SGCModel import SGConvRegressorModel
 from photonai_graph.GraphUtilities import get_random_labels
 
 
@@ -20,10 +20,10 @@ def test_gat_regressor_output_shape(self):
         gat_clf = SGConvRegressorModel(nn_epochs=20)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
 
     def test_gat_regressor_output_hidden_dim(self):
         gat_clf = SGConvRegressorModel(hidden_dim=128)
         gat_clf.fit(self.Xrandom4d, self.y)
         output = gat_clf.predict(self.Xrandom4d)
-        self.assertEqual(output.shape, (20, 1))
+        self.assertEqual(output.shape, self.y.shape)
diff --git a/test/utility_tests/test_save_graphs.py b/test/utility_tests/test_save_graphs.py
deleted file mode 100644
index 7ef3696..0000000
--- a/test/utility_tests/test_save_graphs.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import unittest
-import networkx as nx
-import os
-from photonai_graph.GraphUtilities import save_graphs
-
-
-class SaveGraphsTest(unittest.TestCase):
-
-    def setUp(self):
-        self.graphs = [nx.erdos_renyi_graph(20, 0.3)] * 20
-
-    def test_save(self):
-        directory = "/tmp/"
-        save_graphs(self.graphs, path=directory)
-        self.assertTrue(os.path.exists(directory + "graph_1"), True)
-        os.remove(directory+"graph_1")
-
-    def test_exception(self):
-        with self.assertRaises(NotImplementedError):
-            save_graphs(self.graphs, input_format="dense")