How samuelfriher/tfg-fraud-detection-gnn works
1. This repository is a benchmarking and experimentation framework for AML fraud detection, centered on two competing model families: traditional tabular ML and a custom graph neural network. The main architectural boundary is the pipeline layer, which orchestrates preprocessing, model training, evaluation, explainability, and experiment logging.
2. The entry point is a unified CLI in src/cli.py. It routes commands into separate workflows for:
• traditional model training,
• GNN training,
• GNN hyperparameter grid search,
• results retrieval/export.
This makes the CLI the top-level control plane for the whole repo.
3. Configuration is externalized in TOML files under configs/ and parsed by src/config/experiment_config.py. The presence of separate small/large configs for traditional and GNN runs indicates the pipelines are parameterized for different dataset scales or experiment scopes rather than hard-coded around one setup.
The data subsystem in src/data/ is split by responsibility
• preprocessing of the raw AML dataset,
• graph construction,
• node and edge feature extraction,
• Neo4j-based feature loading/extraction.
This suggests a two-path data flow: a tabular path for traditional models and a graph-enriched path for the GNN.
5. Neo4j is an important architectural dependency, but it is optional at runtime. The README and the neo4j_loader / neo4j_feature_extractor modules indicate the system can project graphs into Neo4j and pull graph analytics features such as PageRank, WCC, and FastRP. If Neo4j is unavailable, the pipeline falls back to basic statistics, so graph feature augmentation is a best-effort enrichment layer rather than a hard requirement.
6. Traditional ML lives in src/models/traditional/ as separate wrappers for XGBoost, LightGBM, Random Forest, and SVM. This is a model-adapter design: the pipeline can treat each algorithm through a common interface while still allowing library-specific behavior, including GPU-accelerated cuML-backed execution for some models when available.
7. The GNN stack is custom and self-contained under src/models/gnn/. It includes model definition, layers, loss, evaluation, and data loading utilities. The README identifies the core architecture as a MEGA-PNA encoder with an edge classifier, which means the model predicts at the transaction-edge level rather than only node classes. This is the main differentiator from the traditional tabular models.
8. Pipeline orchestration is separated from model code in src/pipelines/. There are distinct orchestrators for traditional training, GNN training, and GNN grid search, plus a results exporter. This separation matters because experiment coordination, not model definition, owns cross-cutting concerns like data preparation order, metric aggregation, and artifact generation.
9. Explainability is a first-class subsystem in src/explainability/. Traditional models use SHAP-based explainers, while the GNN path uses GNNExplainer for edge-level interpretation. Visualizers are isolated from explainer logic, which implies the repo treats explanation generation and explanation presentation as separate concerns.
10. Experiment tracking is handled in src/tracking/experiment_tracker.py, with MLflow as the persistence layer and SQLite as the local backend according to the README. The tracking layer sits beside pipelines rather than inside models, so metrics, parameters, and artifacts are captured uniformly regardless of which algorithm family is running.
11. src/utils/ provides supporting infrastructure rather than domain logic: filesystem/path helpers, logging, GPU availability checks, and data management. The GPU check is especially relevant because it influences runtime behavior in both the traditional SVM/Random Forest path and the GNN stack.
12. The repo is designed for reproducible local development through the .devcontainer/ setup, which bundles Python and a Neo4j service. That means the expected runtime architecture is “single developer environment with local services,” not a distributed production deployment.
13. Tests mirror the source tree closely under tests/, which reinforces the intended boundaries: configuration, data preparation, models, pipelines, tracking, and utilities are all independently verifiable. This test layout is useful for diagramming because it reflects the actual module seams in the system.
Main entry points: gnn_pipeline.py GNN train · cli.py CLI · model.py GNN model · traditional_pipeline.py Trad train
Components 30
Control plane
Entry points and configuration for experiment execution.
Data layer
Preprocessing, graph construction, and optional Neo4j enrichment.
- Preprocess · Cleans and prepares raw AML data for downstream paths.
src/data/preprocessor.py - Graph build · Builds graph representations used by the GNN pipeline.
src/data/graph_builder.py - Node feats · Derives node-level features for graph workflows.
src/data/node_feature_extractor.py - Edge feats · Derives edge-level features for graph workflows.
src/data/edge_feature_extractor.py - Neo4j load · Loads graph data into Neo4j when available.
src/data/neo4j_loader.py - Neo4j feats · Pulls optional graph analytics features from Neo4j.
src/data/neo4j_feature_extractor.py
Model layer
Traditional tabular models and the custom GNN stack.
- Tabular ML · Common wrapper layer for traditional classification models.
- XGBoost · Gradient-boosted tree classifier adapter.
src/models/traditional/xgboost_model.py - LightGBM · LightGBM classifier adapter, including accelerated variants when available.
src/models/traditional/lightgbm_model.py - Random Forest · Random forest classifier adapter.
src/models/traditional/random_forest_model.py - SVM · Support vector machine adapter.
src/models/traditional/svm_model.py - GNN model · Custom graph neural network for transaction-edge classification.
src/models/gnn/model.py - GNN layers · Core GNN layers and encoder components.
src/models/gnn/layers.py - GNN loss · Training loss functions for the GNN stack.
src/models/gnn/loss.py - GNN eval · Evaluation and scoring utilities for GNN runs.
src/models/gnn/evaluator.py - GNN data · Builds GNN-ready batches and inputs from graph data.
src/models/gnn/data_loader.py
Pipelines
Orchestrators that run training, search, and export workflows.
- Trad train · Coordinates preprocessing, model fitting, and evaluation for tabular ML.
src/pipelines/traditional_pipeline.py - GNN train · Coordinates graph preparation, training, evaluation, and logging for GNN runs.
src/pipelines/gnn_pipeline.py - Grid search · Runs GNN hyperparameter sweeps over configured search spaces.
src/pipelines/gnn_grid_search.py - Results export · Retrieves and exports experiment results.
src/pipelines/results_exporter.py
Explainability
Model interpretation for both model families.
- SHAP · Explains traditional models with SHAP-based methods.
src/explainability/traditional_explainer.py - GNNExplainer · Explains GNN predictions at the edge level.
src/explainability/gnn_explainer.py - Visuals · Renders explanation outputs for inspection.
src/explainability/visualizers.py
Tracking
Experiment logging and persistence.
- MLflow · Persists metrics, parameters, and artifacts to the tracking backend.
src/tracking/experiment_tracker.py
Utilities
Shared support code for runtime and filesystem concerns.
Developer runtime
Local development environment and service composition.
- Dev env · Local development container with Python and Neo4j services.
.devcontainer/devcontainer.json
Connections 35
- CLI → Config: loads · CLI reads experiment settings to route commands.
- CLI → Trad train: runs · CLI launches traditional training workflows.
- CLI → GNN train: runs · CLI launches GNN training workflows.
- CLI → Grid search: runs · CLI launches GNN hyperparameter search.
- CLI → Results export: exports · CLI exposes results retrieval and export commands.
- Trad train → Preprocess: uses · Tabular pipeline starts from shared preprocessing.
- Trad train → Tabular ML: trains · Pipeline dispatches to the selected traditional model adapter.
- GNN train → Preprocess: uses · GNN pipeline begins with shared preprocessing.
- GNN train → Graph build: builds · Pipeline constructs graph inputs for the GNN.
- GNN train → GNN data: feeds · Pipeline prepares model-ready graph batches.
- GNN train → GNN model: trains · Pipeline trains the custom graph model.
- Grid search → GNN train: iterates · Search wraps repeated GNN training runs.
- GNN train → MLflow: logs · Pipeline records parameters, metrics, and artifacts.
- Trad train → MLflow: logs · Traditional runs are tracked uniformly.
- Tabular ML → XGBoost: includes · Traditional family includes XGBoost.
- Tabular ML → LightGBM: includes · Traditional family includes LightGBM.
- Tabular ML → Random Forest: includes · Traditional family includes Random Forest.
- Tabular ML → SVM: includes · Traditional family includes SVM.
- GNN model → GNN layers: composes · Model is built from custom GNN layers.
- GNN model → GNN loss: optimizes · Training uses the GNN loss module.
- GNN model → GNN eval: evaluates · Evaluation is handled by the GNN evaluator.
- Graph build → Node feats: derives · Graph construction uses node features.
- Graph build → Edge feats: derives · Graph construction uses edge features.
- Graph build → Neo4j load: optional sync · Graphs may be projected into Neo4j.
- Neo4j load → Neo4j feats: enriches · Neo4j data supports graph analytics feature extraction.
- Neo4j feats → Graph build: augments · Optional graph analytics features feed back into graph inputs.
- Tabular ML → SHAP: explains · Traditional models are interpreted with SHAP.
- GNN model → GNNExplainer: explains · GNN predictions are interpreted with GNNExplainer.
- SHAP → Visuals: renders · Traditional explanations are visualized separately.
- GNNExplainer → Visuals: renders · GNN explanations are visualized separately.
- GNN train → GPU check: checks · GPU availability can alter runtime behavior.
- Trad train → GPU check: checks · GPU availability can affect some traditional backends.
- CLI → Paths: resolves · Shared paths support command execution.
- CLI → Logger: logs · CLI and workflows share logging setup.
- Dev env → Neo4j load: provides · Local container includes Neo4j for graph enrichment.