Knowledge Graph Intrusion Detection

Build a GNN that detects network intrusions by learning graph structure.

Introduction

30 Second Summary

Traditional security systems look at each network connection in isolation, like checking individual puzzle pieces without seeing the full picture. Real attackers exploit relationships between systems, moving laterally through networks in patterns that per-connection analysis completely misses.

In this project, you will build an intrusion detection system that models network traffic as a knowledge graph and uses a Graph Neural Network to catch attacks that traditional classifiers miss. You will also predict which services are likely to be targeted next based on graph structure.

What You'll Build

Run your finished system against network traffic data and watch it flag rare attack types that a traditional Random Forest completely overlooks, then query it to see a ranked threat forecast of services likely to be targeted next.

By the end of this project, you'll have:

  • A knowledge graph built from real network traffic where you can visualize how attack patterns cluster around specific services differently from normal traffic.
  • A trained GraphSAGE model that detects rare intrusions (R2L and U2R) where your Random Forest baseline falls short, proving graph structure adds detection signal flat features cannot.
  • A threat prediction system you can query to see which services are most likely to be attacked next, ranked by graph centrality and learned embeddings.
  • Secret Mission: Enrich the knowledge graph with MITRE ATT&CK technique mappings to classify detected intrusions by attack tactic.

Are there any prerequisites?

This project assumes Python proficiency and basic ML knowledge (classification, train/test splits, evaluation metrics). You do not need prior experience with graph neural networks or knowledge graphs.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Traditional intrusion detection treats every network flow as an independent row in a spreadsheet. You're about to build something that sees the bigger picture by modeling traffic as a knowledge graph, then using a Graph Neural Network to catch the attacks that flat classifiers miss.

Set Up Your Environment

Building a graph-based intrusion detection system requires specialized libraries that work together. PyTorch handles the neural network training, while PyTorch Geometric adds graph-specific layers on top.

Getting version alignment right between PyTorch and PyTorch Geometric avoids painful debugging later. In this step, you'll install everything and download the network traffic dataset you'll be working with.

In this step, get ready to:
  • Install PyTorch and PyTorch Geometric with compatible versions.
  • Install data science libraries (NetworkX, scikit-learn, pandas, matplotlib).
  • Download the NSL-KDD dataset and create your project notebook.
Check Python version and install PyTorch CPU

PyTorch Geometric requires Python 3.11 or higher. You'll confirm your version first, then install PyTorch configured for CPU training.

  • Press the Windows key to open your search bar.
  • Type Terminal and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir intrusion-detection
cd intrusion-detection
  • Verify your Python version by running:
python --version

You should see Python 3.11 or higher printed.

  • Install PyTorch (CPU version) by running this command:
pip install torch --index-url https://download.pytorch.org/whl/cpu

Why the CPU version?

The knowledge graph you'll build in this project has around 1,000-3,000 nodes. That's small enough to train on CPU in seconds. The CPU-only install is much smaller and avoids GPU driver complications.

  • Confirm PyTorch installed correctly by running:
python -c "import torch; print(torch.__version__)"

You should see a version number printed (e.g. 2.6.0+cpu). Take note of the number before +cpu because you'll need it in the next substep.

Seeing an error during install?

Make sure you typed the full URL after --index-url exactly as shown. A missing /cpu at the end will cause the download to fail.

If you see a permissions error, try running the command with --user appended to the end.

Help me debug my PyTorch CPU installation on Windows.

Install PyTorch Geometric and extensions

PyTorch Geometric adds graph-specific neural network layers (like GraphSAGE) on top of PyTorch. Its optional extension packages accelerate sparse matrix operations that graph computations rely on.

  • Install the core PyTorch Geometric package by running:
pip install torch-geometric
  • Install the accelerated extension packages by running the command below, replacing the version number with the PyTorch version you noted earlier (the part before +cpu):
pip install pyg-lib torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-[[TORCH_VERSION="your PyTorch version e.g. 2.6.0"]]+cpu.html

Why does the URL need my PyTorch version?

The extension packages are compiled against specific PyTorch versions. The URL points pip at pre-built wheels that match your exact PyTorch install. A mismatched version causes cryptic import errors at runtime.

  • Confirm PyTorch Geometric installed correctly by running:
python -c "import torch_geometric; print(torch_geometric.__version__)"

You should see a version number printed without any errors.

Import fails with a missing module error?

Double-check the version in your extensions URL matches your PyTorch version exactly. Run python -c "import torch; print(torch.__version__)" again and compare.

If you see a network timeout, retry the install command. The PyG wheel server can be slow on first download.

Help me fix my PyTorch Geometric installation.

Install remaining dependencies and download the dataset

You still need NetworkX for building the graph structure, scikit-learn for the baseline classifier, and pandas/matplotlib for data handling and visualization.

  • Install the remaining data science packages by running:
pip install networkx scikit-learn pandas matplotlib seaborn
  • Confirm all packages import correctly by running:
python -c "import networkx, sklearn, pandas, matplotlib; print('All packages ready')"

You should see All packages ready printed with no errors.

Now you need the NSL-KDD dataset. This is a benchmark network intrusion detection dataset with 41 features per network flow and 5 attack categories.

  • Open this URL in your browser:
https://www.kaggle.com/datasets/hassan06/nslkdd
  • Sign in to Kaggle (or create a free account if you don't have one).
  • Click the Download button to download the dataset as a zip file.
  • Extract the zip file contents into your intrusion-detection folder on the Desktop.

What files should I see?

The dataset contains KDDTrain+.TXT and KDDTest+.TXT as the main files you'll use. Despite the .TXT extension, these are CSV-formatted with comma-separated values. You'll load them directly with pandas in the next step.

No Kaggle account?

The same dataset is available as ready-made CSV files on HuggingFace at huggingface.co/datasets/Mireu-Lab/NSL-KDD. Download train.csv and test.csv from the Files tab and place them in your intrusion-detection folder.

  • Confirm the dataset files are in your project folder by running:
ls

You should see KDDTrain+.TXT and KDDTest+.TXT listed (or train.csv and test.csv if you used HuggingFace).

Finally, create the Jupyter notebook you'll work in for the rest of this project.

  • Switch back to VS Code.
  • Click File in the top menu bar.
  • Click Open Folder.
  • Select the intrusion-detection folder on your Desktop, then click Select Folder.
  • Press Ctrl+Shift+P to open the Command Palette.
  • Type Create: New Jupyter Notebook and select it from the dropdown.
  • Save the notebook immediately with Ctrl+S.
  • Name the file intrusion_detection.ipynb and click Save.

You should see intrusion_detection.ipynb open as a blank notebook in VS Code with an empty code cell ready.

Before you run the final check, what do you think will happen when you import both PyTorch and PyTorch Geometric in the same command?

  • Run the final verification back in your terminal:
python -c "import torch; import torch_geometric; print(torch.__version__, torch_geometric.__version__)"

You should see both version numbers printed side by side with no errors. This confirms PyTorch and PyTorch Geometric are installed and compatible with each other.

Seeing a version mismatch or import error?

If PyTorch imports fine but torch_geometric fails, uninstall and reinstall the extensions with the correct version URL. Run python -c "import torch; print(torch.__version__)" to get your exact version and make sure the URL matches.

Help me resolve a PyTorch Geometric version compatibility issue.

Your environment is ready. Next up, you'll load the NSL-KDD data into your notebook, train a Random Forest baseline, and see exactly where traditional classifiers fall short on rare attack types.

Load Data and Build the Baseline Classifier

Your environment is ready with all the libraries installed. Now you need to see what a traditional machine learning classifier can actually achieve on network intrusion data.

Before building anything graph-based, you need a baseline. A Random Forest classifier will show you the performance ceiling for non-graph methods. More importantly, it will reveal a critical blind spot that flat classifiers have when dealing with rare attack types.

In this step, get ready to:
  • Load and preprocess the NSL-KDD dataset into a pandas DataFrame.
  • Train a Random Forest baseline and evaluate its per-class performance.
  • Identify where the baseline fails and understand why.
Load and preprocess NSL-KDD

The NSL-KDD dataset contains network traffic records. Each row represents a single network connection (called a flow) with 41 features describing that connection's properties.

  • In your intrusion_detection.ipynb notebook, create a new cell and import the required libraries by pasting this code:
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from sklearn.preprocessing import LabelEncoder

# NSL-KDD column names (41 features + label + difficulty)
columns = ['duration', 'protocol_type', 'service', 'flag', 'src_bytes',
    'dst_bytes', 'land', 'wrong_fragment', 'urgent', 'hot',
    'num_failed_logins', 'logged_in', 'num_compromised', 'root_shell',
    'su_attempted', 'num_root', 'num_file_creations', 'num_shells',
    'num_access_files', 'num_outbound_cmds', 'is_host_login',
    'is_guest_login', 'count', 'srv_count', 'serror_rate',
    'srv_serror_rate', 'rerror_rate', 'srv_rerror_rate', 'same_srv_rate',
    'diff_srv_rate', 'srv_diff_host_rate', 'dst_host_count',
    'dst_host_srv_count', 'dst_host_same_srv_rate',
    'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
    'dst_host_srv_diff_host_rate', 'dst_host_serror_rate',
    'dst_host_srv_serror_rate', 'dst_host_rerror_rate',
    'dst_host_srv_rerror_rate', 'label', 'difficulty_level']

# Load the NSL-KDD training dataset
df = pd.read_csv('KDDTrain+.txt', header=None, names=columns)
print(f"Dataset loaded: {df.shape[0]} records, {df.shape[1]} columns")
print(df.head(3))

What does this code do?

The KDDTrain+.txt file has no column headers, so you define the 43 column names explicitly. The first 41 are network flow features (like src_bytes, duration, service). The last two are label (the specific attack name) and difficulty_level (how hard the record is to classify).

Three of the 41 features are categorical (text) rather than numeric: protocol_type, service, and flag. You will encode these next.

  • Run the cell.
  • You should see output showing approximately 125,000 records across 43 columns, with a preview of the first three rows.

Seeing a FileNotFoundError?

Make sure the KDDTrain+.txt file is in the same folder as your notebook. If you downloaded it to a different location, update the file path in pd.read_csv() to point to where the file actually lives.

help me fix my file path for loading the NSL-KDD dataset

Now you need to convert the three categorical columns into numbers and group the 39 specific attack labels into 5 broader categories.

  • Create a new cell and paste this preprocessing code:
# Label encode the three categorical features
le_protocol = LabelEncoder()
le_service = LabelEncoder()
le_flag = LabelEncoder()

df['protocol_type'] = le_protocol.fit_transform(df['protocol_type'])
df['service'] = le_service.fit_transform(df['service'])
df['flag'] = le_flag.fit_transform(df['flag'])

# Map 39 specific attack labels to 5 categories
attack_mapping = {
    'normal': 'Normal',
    'back': 'DoS', 'land': 'DoS', 'neptune': 'DoS', 'pod': 'DoS',
    'smurf': 'DoS', 'teardrop': 'DoS', 'apache2': 'DoS',
    'udpstorm': 'DoS', 'processtable': 'DoS', 'mailbomb': 'DoS',
    'satan': 'Probe', 'ipsweep': 'Probe', 'nmap': 'Probe',
    'portsweep': 'Probe', 'mscan': 'Probe', 'saint': 'Probe',
    'guess_passwd': 'R2L', 'ftp_write': 'R2L', 'imap': 'R2L',
    'phf': 'R2L', 'multihop': 'R2L', 'warezmaster': 'R2L',
    'warezclient': 'R2L', 'spy': 'R2L', 'xlock': 'R2L',
    'xsnoop': 'R2L', 'snmpguess': 'R2L', 'snmpgetattack': 'R2L',
    'httptunnel': 'R2L', 'sendmail': 'R2L', 'named': 'R2L',
    'buffer_overflow': 'U2R', 'loadmodule': 'U2R', 'rootkit': 'U2R',
    'perl': 'U2R', 'sqlattack': 'U2R', 'xterm': 'U2R', 'ps': 'U2R'
}
df['attack_category'] = df['label'].map(attack_mapping)
print(f"Attack categories: {df['attack_category'].unique()}")
print(df['attack_category'].value_counts())

What does this code do?

The LabelEncoder converts each categorical column into integers. The Random Forest needs numeric inputs, so tcp becomes 0, udp becomes 1, etc.

The attack_mapping dictionary groups 39 specific attack names (like neptune, smurf, buffer_overflow) into 5 high-level categories: Normal, DoS, Probe, R2L, and U2R. This gives you a cleaner classification target.

  • Run the cell.
  • You should see the five attack categories listed, along with a count showing that Normal and DoS dominate the dataset while R2L and U2R have far fewer samples.

Next you will split the data and train the baseline. Create a new cell with the train/test split:

  • Create a new cell and prepare the features and train/test split:
# Prepare features (41 numeric columns) and target labels
feature_cols = [c for c in df.columns
    if c not in ['label', 'difficulty_level', 'attack_category']]
X = df[feature_cols].values
y = df['attack_category'].values

# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)
print(f"Train size: {len(X_train)}, Test size: {len(X_test)}")
print(f"Features per sample: {X_train.shape[1]}")

Why stratify the split?

The stratify=y parameter ensures the train and test sets have the same proportion of each attack category. Without it, the rare categories (R2L, U2R) might end up underrepresented in one set, making evaluation unreliable.

  • Run the cell. You should see roughly 87,000 training samples and 38,000 test samples, each with 41 features.
Train the Random Forest baseline

With the data split, you can now train a classifier and see how well a traditional approach handles intrusion detection.

  • Create a new cell and train the Random Forest model:
# Train a Random Forest classifier on flat features
rf_model = RandomForestClassifier(
    n_estimators=100, random_state=42, n_jobs=-1
)
rf_model.fit(X_train, y_train)
print("Random Forest training complete.")
  • Run the cell. Training should complete in a few seconds.

Before you look at the results, think about this: the Random Forest sees each network flow as an independent row of 41 numbers. Do you think it will perform equally well on all five attack categories? Or will some be harder to detect?

  • Create a new cell and generate the classification report:
# Generate predictions and evaluate per-class performance
y_pred = rf_model.predict(X_test)
print("Random Forest Classification Report:")
print("=" * 60)
print(classification_report(y_test, y_pred))
  • Run the cell and study the output carefully.

You will see that the overall accuracy exceeds 95%. The model handles Normal, DoS, and Probe categories well. But look at the recall for R2L and U2R. Both fall below 60%. The model is missing the majority of these rare attacks.

What is recall telling you here?

Recall measures what fraction of actual attacks in a category were correctly detected. A recall of 0.55 on R2L means the model only catches 55% of remote-to-local attacks. The other 45% slip through undetected. In a real security operations center, those missed attacks could mean a full breach.

Analyze the failure

The classification report showed the gap. Now you need to understand why it exists.

  • Create a new cell and print the class distribution to quantify the imbalance:
# Show class distribution to reveal the imbalance
print("Class Distribution in Training Data:")
print("-" * 40)
for category in ['Normal', 'DoS', 'Probe', 'R2L', 'U2R']:
    count = (y_train == category).sum()
    pct = count / len(y_train) * 100
    print(f"  {category:8s}: {count:6d} samples ({pct:.1f}%)")

print("\nThe Random Forest treats each network flow independently.")
print("It cannot see that R2L/U2R attacks share infrastructure")
print("patterns with other flows - patterns only visible")
print("through connections between hosts and services.")

Why does the baseline fail on rare attacks?

Two problems compound here. First, R2L and U2R together represent less than 3% of the training data. The model has very few examples to learn from.

Second, these attacks look similar to normal traffic when you examine each flow in isolation. An R2L attack might use the same service and byte counts as a legitimate connection. The distinguishing signal is relational: multiple R2L attacks target the same service from different sources, forming a pattern visible only when you connect the flows through their shared infrastructure.

  • Run the cell. You should see Normal and DoS dominating (each above 40%), Probe around 9%, and R2L and U2R together representing less than 3% of all samples.

Getting different percentages?

If your numbers differ slightly from expectations, confirm you used random_state=42 in train_test_split. Also confirm your attack_mapping dictionary covers all attack labels in your dataset. Any unmapped labels will produce NaN values in the attack_category column.

help me debug my attack category mapping

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all five cells present and running without errors.

ⓧ I'd like to double check the full code

Scroll up through your notebook and confirm these five cells exist in order:

  • Cell 1: Imports + column names + pd.read_csv loading df.
  • Cell 2: Label encoding of protocol_type, service, flag + attack mapping creating attack_category.
  • Cell 3: Feature extraction + train_test_split into X_train, X_test, y_train, y_test.
  • Cell 4: rf_model training + classification_report output.
  • Cell 5: Class distribution loop showing per-category percentages.

You now have hard evidence that a traditional classifier falls short. It achieves impressive overall accuracy, but misses the attacks that matter most. In the next step, you will build a knowledge graph from this same data, creating the structural connections that make rare attacks visible.

Construct the Knowledge Graph

Your Random Forest achieved over 95% overall accuracy, but it exposed a critical blind spot. Rare attack types like R2L and U2R slipped through because the classifier treats every network flow as an isolated row of numbers.

In this step, you will construct a knowledge graph that models network traffic as a web of relationships. Hosts, services, and protocols become nodes. Connection patterns become edges. This relational structure gives your model the context it needs to spot attack campaigns that span multiple flows.

In this step, get ready to:
  • Design a graph ontology with three node types and two edge types.
  • Build a directed knowledge graph from the NSL-KDD DataFrame using NetworkX.
  • Visualize attack clustering patterns that flat classifiers cannot detect.
Design the ontology

An ontology defines the vocabulary of your knowledge graph. It specifies what types of entities exist and how they relate to each other.

Your NSL-KDD dataset has three natural entity types. Source hosts produce traffic, destination services receive it, and protocols carry it. Two relationship types connect them: connects_to links hosts to services, and uses_protocol links hosts to protocols.

  • Create a new cell in your notebook and define the ontology schema by running this code:
# Knowledge Graph Ontology Schema
# Defines entity types and relationships for network traffic modeling
ontology = {
    'node_types': {
        'source_host': 'Hosts grouped by dst_host_count traffic patterns',
        'service': 'Destination services from service column',
        'protocol': 'Network protocols from protocol_type column'
    },
    'edge_types': {
        'connects_to': 'source_host -> service (carries flow statistics)',
        'uses_protocol': 'source_host -> protocol (carries connection metadata)'
    },
    'node_features': ['src_bytes', 'dst_bytes', 'serror_rate',
                      'dst_host_count', 'count']
}

print("Ontology defined:")
print(f"  Node types: {list(ontology['node_types'].keys())}")
print(f"  Edge types: {list(ontology['edge_types'].keys())}")
print(f"  Node features: {ontology['node_features']}")

What does this code do?

This dictionary is a blueprint for the graph you are about to build. It declares the three node types (source hosts, services, protocols) and two edge types (connects_to, uses_protocol).

The node_features list names the flow statistics each host node will carry. These become the input features for the GNN later.

  • Run the cell.

You should see the three node types, two edge types, and five node features printed out.

Why group hosts by dst_host_count?

The NSL-KDD dataset does not include raw IP addresses. Instead, dst_host_count measures how many connections targeted the same host in the last 100 connections. This creates a natural grouping that captures traffic volume profiles without needing IPs.

Build the graph nodes from flows

With the ontology defined, you can now instantiate the actual graph. Each unique value in the service, protocol_type, and dst_host_count columns becomes a node in the graph.

  • Create a new cell and build the graph nodes by running this code:
import networkx as nx

G = nx.DiGraph()

# Create protocol nodes from unique protocols
for proto_id in df['protocol_type'].unique():
    G.add_node(f"proto_{proto_id}", node_type='protocol')

# Create service nodes from unique services
for service_id in df['service'].unique():
    G.add_node(f"service_{service_id}", node_type='service')

# Create source host nodes grouped by traffic volume patterns
for host_id, group in df.groupby('dst_host_count'):
    G.add_node(f"host_{host_id}", node_type='source_host',
               avg_src_bytes=float(group['src_bytes'].mean()),
               avg_dst_bytes=float(group['dst_bytes'].mean()),
               serror_rate=float(group['serror_rate'].mean()),
               connection_count=len(group))

print(f"Nodes created: {G.number_of_nodes()}")

What does this code do?

  • Creates a directed graph using NetworkX's DiGraph class.
  • Adds protocol nodes (3 unique values: TCP, UDP, ICMP encoded as integers).
  • Adds service nodes (around 70 unique values like HTTP, FTP, SMTP).
  • Adds host nodes (up to 256 groups from dst_host_count), each carrying aggregated flow statistics as node features.
  • Run the cell.

You should see a count of a few hundred nodes. These are all the entities in your knowledge graph, ready to be connected.

Seeing a NameError for df?

Make sure you ran all cells from the previous step first. The DataFrame df must exist from loading the NSL-KDD data in Step 2.

Help me debug a NameError in my knowledge graph construction

Add edges to connect the graph

Nodes alone are just isolated points. Edges encode the relationships that make this a knowledge graph. Each edge carries metadata about the connection pattern it represents.

  • Create a new cell and add the graph edges by running this code:
# Add connects_to edges (host -> service)
for host_id, group in df.groupby('dst_host_count'):
    for service_id in group['service'].unique():
        subset = group[group['service'] == service_id]
        G.add_edge(f"host_{host_id}", f"service_{service_id}",
                   edge_type='connects_to',
                   flow_count=len(subset),
                   avg_src_bytes=float(subset['src_bytes'].mean()))

# Add uses_protocol edges (host -> protocol)
for host_id, group in df.groupby('dst_host_count'):
    for proto_id in group['protocol_type'].unique():
        subset = group[group['protocol_type'] == proto_id]
        G.add_edge(f"host_{host_id}", f"proto_{proto_id}",
                   edge_type='uses_protocol',
                   connection_count=len(subset))

print(f"Knowledge Graph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

What does this code do?

  • Creates connects_to edges from each host group to every service it contacted, carrying flow count and average byte volume.
  • Creates uses_protocol edges from each host group to the protocols it used, carrying connection counts.
  • These edge attributes become the relational signal the GNN will learn from. A DoS host connecting to one service over thousands of flows looks very different from an R2L host making a handful of quiet connections.

Before you run this cell, how many edges do you expect? Consider that each host group connects to multiple services and protocols.

  • Run the cell.

You should see hundreds of nodes and thousands of edges. The graph is now a rich relational structure encoding how hosts interact with services through protocols.

Graph has 0 edges?

Check that you ran the node-creation cell first. If you ran the edge cell without building nodes, the graph object G may have been re-initialized. Rerun both cells in order.

Help me debug why my knowledge graph has no edges

Visualize the knowledge graph

A graph with thousands of edges is too dense to plot in full. By sampling a subgraph and coloring nodes by their dominant attack category, you can see how attack patterns cluster differently from normal traffic.

  • Create a new cell and plot the knowledge graph subgraph by running this code:
import matplotlib.pyplot as plt

# Determine dominant attack category per host group
host_attack_map = {}
for host_id, group in df.groupby('dst_host_count'):
    host_attack_map[f"host_{host_id}"] = group['attack_category'].mode().iloc[0]

# Sample a subgraph for readable visualization
host_nodes = [n for n in G.nodes() if G.nodes[n].get('node_type') == 'source_host']
subgraph_nodes = set(host_nodes[:30])
for host in list(subgraph_nodes):
    subgraph_nodes.update(G.successors(host))
sub_G = G.subgraph(subgraph_nodes)

# Color nodes by attack category
category_colors = {'Normal': '#2ecc71', 'DoS': '#e74c3c',
                   'Probe': '#f39c12', 'R2L': '#9b59b6', 'U2R': '#1abc9c'}
color_map = []
for node in sub_G.nodes():
    cat = host_attack_map.get(node, None)
    color_map.append(category_colors.get(cat, '#95a5a6'))

# Plot the knowledge graph subgraph
pos = nx.spring_layout(sub_G, seed=42, k=0.3)
plt.figure(figsize=(12, 8))
nx.draw(sub_G, pos, node_color=color_map, node_size=80,
        edge_color='lightgray', alpha=0.7, with_labels=False)
plt.title("Knowledge Graph: DoS Clustering vs R2L Low-Volume Patterns")
plt.tight_layout()
plt.show()

What does this code do?

  • Determines the dominant attack category for each host group by finding the most common label in that group.
  • Samples 30 host nodes plus all their connected service and protocol nodes to create a readable subgraph.
  • Colors each host node by its dominant attack category: green for Normal, red for DoS, orange for Probe, purple for R2L, teal for U2R. Service and protocol nodes appear gray.
  • Uses a spring layout that pushes densely connected clusters together, making attack patterns visually distinct.
  • Run the cell.

You should see a graph visualization where red DoS nodes cluster densely around shared service nodes (high fan-in), while purple R2L nodes create sparse, low-volume connections to specific services.

Why do attack types cluster differently?

DoS attacks flood a target service with thousands of connections. In the graph, this creates dense clusters of host nodes all pointing at the same service. R2L attacks use legitimate-looking connections to a few services. Their graph footprint is sparse and scattered.

This structural difference is invisible to a flat classifier that only sees one row at a time. The GNN you build next will learn to propagate information along these edges, detecting attacks by their neighborhood shape.

Plot is blank or too cluttered?

If the plot is blank, confirm that sub_G has nodes by running print(sub_G.number_of_nodes()). If it returns 0, re-run the node and edge cells above first.

If the plot is too cluttered, try changing host_nodes[:30] to host_nodes[:15] to sample fewer nodes.

Help me fix my knowledge graph visualization

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all four new cells (ontology, nodes, edges, visualization).

ⓧ I'd like to double check the full code

Scroll up and confirm you have these four cells in order, each producing output:

  • The ontology schema cell printing node types, edge types, and features.
  • The node creation cell printing the node count.
  • The edge creation cell printing the full graph stats (nodes and edges).
  • The visualization cell showing a colored subgraph plot.

Your knowledge graph now encodes the relational structure that flat classifiers ignore. Next up, you will convert this graph into a format PyTorch Geometric can train on and build the GraphSAGE model that learns from these structural patterns.

Train the GraphSAGE Intrusion Detector

Your knowledge graph now encodes the relational structure that flat classifiers ignore. The Random Forest treated each flow independently, missing multi-hop patterns that connect rare attacks to shared infrastructure.

In this step, you will train a Graph Neural Network that propagates information along edges. GraphSAGE learns to classify nodes by aggregating features from their graph neighborhood, capturing the attack patterns the Random Forest missed entirely.

In this step, get ready to:
  • Convert the NetworkX knowledge graph into a PyTorch Geometric Data object with feature tensors and train/test masks.
  • Define and train a 2-layer GraphSAGE model for node classification.
  • Evaluate the GNN against the Random Forest baseline and compare per-class F1-scores.
Convert the knowledge graph to PyTorch Geometric format

PyTorch Geometric expects graph data as tensors, not as a NetworkX object. You need to extract node features into a tensor x, build an edge_index tensor mapping connections, and create label and mask tensors for supervised training.

  • In your notebook, create a new cell and import PyTorch and PyTorch Geometric by running this code:
import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import SAGEConv
from torch_geometric.utils import from_networkx

# Convert NetworkX graph to PyG Data, grouping all node attrs into x
data = from_networkx(G, group_node_attrs='all')
print(f"Converted graph: {data}")

What does this code do?

The from_networkx utility handles the heavy lifting of converting your NetworkX graph into PyTorch Geometric format.

Setting group_node_attrs='all' tells PyG to concatenate all numeric node attributes into a single feature matrix data.x. It also builds the edge_index tensor automatically from your graph edges.

  • Run the cell. You should see a Data object printed with the edge_index shape and x shape reflecting your graph dimensions.

The converted Data object has features but no labels or train/test split yet. You need to attach the attack category labels and create masks that tell the model which nodes to train on versus evaluate.

  • Create a new cell to add labels and masks by running this code:
# Build label tensor from node 'label' attribute
node_list = list(G.nodes())
node_labels = [G.nodes[node]['label'] for node in node_list]
y = torch.tensor(node_labels, dtype=torch.long)

# Create randomized train/test split masks (80/20)
num_nodes = data.x.size(0)
perm = torch.randperm(num_nodes)
train_size = int(0.8 * num_nodes)
train_mask = torch.zeros(num_nodes, dtype=torch.bool)
test_mask = torch.zeros(num_nodes, dtype=torch.bool)
train_mask[perm[:train_size]] = True
test_mask[perm[train_size:]] = True

# Attach labels and masks to the Data object
data.y = y
data.train_mask = train_mask
data.test_mask = test_mask
print(f"Labels: {y.shape}, Train nodes: {train_mask.sum()}, Test nodes: {test_mask.sum()}")

What does this code do?

The y tensor holds the integer attack category for each node (0=Normal, 1=DoS, 2=Probe, 3=R2L, 4=U2R).

The train_mask and test_mask are boolean tensors that select which nodes the model learns from during training, and which are held out for evaluation. Using torch.randperm ensures the split is randomized rather than sequential.

  • Run the cell. You should see the label tensor shape matching your node count, with roughly 80% of nodes in the training set and 20% in the test set.

Seeing a ValueError about node attributes?

The from_networkx function requires all nodes to have the same set of attributes. If you see this error, check that every node in your graph G has identical attribute keys. Go back to your graph construction cell and verify you are attaching the same features to every node.

help me debug a ValueError from from_networkx about inconsistent node attributes

Define and train the GraphSAGE model

GraphSAGE works by sampling and aggregating features from a node's neighbors at each layer. A 2-layer model means each node's classification is informed by features up to 2 hops away in the graph.

  • Create a new cell to define the GraphSAGE model architecture by running this code:
class GraphSAGE(torch.nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        # First layer: input features to 64-dim hidden representation
        self.conv1 = SAGEConv(in_channels, 64)
        # Second layer: 64-dim hidden to output class predictions
        self.conv2 = SAGEConv(64, num_classes)

    def forward(self, x, edge_index):
        # Aggregate neighbor features, apply ReLU activation
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        # Final layer outputs class logits
        x = self.conv2(x, edge_index)
        return x

# Instantiate the model
num_features = data.x.size(1)
num_classes = len(torch.unique(data.y))
gnn_model = GraphSAGE(num_features, num_classes)
print(gnn_model)

What does this architecture do?

Each SAGEConv layer aggregates features from neighboring nodes using mean aggregation. The first layer compresses the input features down to a 64-dimensional hidden representation.

The second layer maps that 64-dimensional space to the number of attack categories (5 classes). Dropout during training prevents the model from memorizing specific patterns and helps it generalize.

  • Run the cell. You should see the model architecture printed, showing two SAGEConv layers with the input and output dimensions matching your data.

Now you will train the model for 200 epochs using cross-entropy loss. The optimizer adjusts the model weights to minimize classification error on the training nodes.

  • Create a new cell to train the model by running this code:
# Set up optimizer and train the model
optimizer = torch.optim.Adam(gnn_model.parameters(), lr=0.01)
gnn_model.train()

for epoch in range(200):
    optimizer.zero_grad()
    out = gnn_model(data.x, data.edge_index)
    # Compute loss only on training nodes
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

    # Print progress every 50 epochs
    if (epoch + 1) % 50 == 0:
        pred = out.argmax(dim=1)
        train_acc = (pred[data.train_mask] == data.y[data.train_mask]).sum() / data.train_mask.sum()
        print(f"Epoch {epoch+1}/200, Loss: {loss:.4f}, Train Acc: {train_acc:.4f}")

print("Training complete!")

What is happening during training?

Each epoch passes all node features through the model, but loss is only calculated on train_mask nodes. This is transductive learning: the model sees the full graph structure but only receives supervision from labeled training nodes.

Information from unlabeled test nodes still flows through the graph during message passing. This is why GNNs can leverage structural context that per-sample classifiers cannot access.

  • Run the cell. You should see loss decreasing and accuracy increasing across the four progress prints, finishing with a training accuracy well above 0.9.

Training accuracy stuck near random chance?

If accuracy stays around 0.2 (1/5 classes), your node features may not be numeric. Verify that data.x.dtype is torch.float32 by printing it. Non-numeric attributes passed to from_networkx will cause silent failures.

help me debug why my GraphSAGE training accuracy is stuck at random chance

Evaluate and compare against the baseline

The real test is whether graph structure gives the GNN an edge on the attack types the Random Forest struggled with. R2L and U2R are the categories where flat classifiers fail because the attacks look normal in isolation but form distinctive patterns in the graph.

  • Create a new cell to evaluate the GNN and print its classification report by running this code:
from sklearn.metrics import classification_report

# Evaluate on test nodes
gnn_model.eval()
with torch.no_grad():
    out = gnn_model(data.x, data.edge_index)
    pred = out.argmax(dim=1)

# Get predictions and true labels for test set
gnn_test_pred = pred[data.test_mask].numpy()
gnn_test_true = data.y[data.test_mask].numpy()

target_names = ['Normal', 'DoS', 'Probe', 'R2L', 'U2R']
print("=== GraphSAGE Classification Report ===")
print(classification_report(gnn_test_true, gnn_test_pred,
                            target_names=target_names))

What does this code do?

Setting the model to eval() mode disables dropout, ensuring deterministic predictions. The torch.no_grad() context avoids tracking gradients during inference, saving memory.

The classification report shows precision, recall, and F1-score per class. Focus on the R2L and U2R rows to see how the GNN compares to what the Random Forest achieved on those same categories.

  • Run the cell. You should see a full classification report with per-class metrics for all five attack categories.

Now for the moment of truth. Before you run the comparison, think about what you expect: do you think the GNN will beat the Random Forest across all categories, or only on the rare attack types?

  • Create a new cell to generate the side-by-side comparison and bar chart by running this code:
from sklearn.metrics import f1_score
import numpy as np

# Per-class F1 scores for both models
gnn_f1 = f1_score(gnn_test_true, gnn_test_pred, average=None)
rf_f1 = f1_score(y_test, rf_model.predict(X_test), average=None)

# Side-by-side comparison table
print("\n=== F1-Score Comparison: Random Forest vs GraphSAGE ===")
for i, name in enumerate(target_names):
    print(f"{name:10s} | RF: {rf_f1[i]:.3f} | GNN: {gnn_f1[i]:.3f}")

# Bar chart visualization
x_pos = np.arange(len(target_names))
width = 0.35
fig, ax = plt.subplots(figsize=(10, 6))
ax.bar(x_pos - width/2, rf_f1, width, label='Random Forest')
ax.bar(x_pos + width/2, gnn_f1, width, label='GraphSAGE')
ax.set_xlabel('Attack Category')
ax.set_ylabel('F1-Score')
ax.set_title('Per-Class F1-Score: Random Forest vs GraphSAGE')
ax.set_xticks(x_pos)
ax.set_xticklabels(target_names)
ax.legend()
plt.tight_layout()
plt.show()

What does this comparison reveal?

The bar chart makes the structural advantage of graph-based learning visible. On common attack types (Normal, DoS), both models perform similarly because volume alone provides enough signal.

The gap widens on R2L and U2R. These rare attacks look like normal traffic when examined in isolation, but they share service targets and connection patterns with other attacks in the graph. The GNN leverages those multi-hop relationships to correctly classify flows the Random Forest misses.

  • Run the cell. You should see a printed comparison table showing higher F1-scores for the GNN on R2L and U2R, and a bar chart clearly visualizing the improvement on rare attack categories.

GNN scores look similar to Random Forest?

If the GNN is not outperforming on R2L/U2R, your graph may not encode enough relational structure. Check that your knowledge graph construction in the previous step created edges between nodes that share services or protocols. A graph with few edges gives the GNN nothing to propagate.

help me understand why my GNN is not outperforming the Random Forest on rare attack types

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all cells run.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm these cells are present in order:

  • Imports cell with torch, torch.nn.functional, Data, SAGEConv, and from_networkx.
  • Conversion cell using from_networkx(G, group_node_attrs='all').
  • Labels and masks cell attaching y, train_mask, test_mask to data.
  • Model definition cell with GraphSAGE class and gnn_model instantiation.
  • Training loop cell running 200 epochs with Adam optimizer.
  • Evaluation cell printing the GNN classification report.
  • Comparison cell with side-by-side F1 table and bar chart.

Your GraphSAGE model now demonstrates that graph structure provides real detection signal for rare attacks. Next up, you will use the knowledge graph's structure for threat prediction, identifying which services are likely to be targeted based on emerging patterns.

Build the Threat Prediction System

Your GraphSAGE model can now detect intrusions more accurately than the Random Forest baseline, especially on rare attack types like R2L and U2R. That answers "what happened?" But security teams need to answer a harder question: "what will happen next?"

In this step, you'll analyze the knowledge graph's structure to predict which services are likely to be targeted. By combining graph centrality metrics with the GNN's learned embeddings, you'll build a threat scoring system that identifies future attack targets from structural patterns alone.

In this step, get ready to:
  • Compute centrality metrics (PageRank, betweenness, in-degree) for every node in the knowledge graph.
  • Build a threat prediction function that scores services by similarity to known attack targets.
  • Validate predictions against held-out test data to confirm predictive capability.
Compute graph centrality metrics

Centrality metrics quantify how structurally important each node is in the graph. A service with high PageRank receives many connections from other important services, making it a high-value target.

  • Add a new cell in your notebook and compute three centrality measures for every node by running this code:
# Compute PageRank to identify high-value nodes
pagerank_scores = nx.pagerank(G)

# Compute betweenness centrality to find attack stepping stones
betweenness_scores = nx.betweenness_centrality(G)

# Compute in-degree to measure direct incoming connections
indegree_scores = dict(G.in_degree())

# Display top 10 nodes by PageRank
top_pagerank = sorted(pagerank_scores.items(), key=lambda x: x[1], reverse=True)[:10]
print("Top 10 nodes by PageRank:")
for node, score in top_pagerank:
    print(f"  {node}: {score:.4f}")

What do these metrics tell you?

  • PageRank scores nodes by how many other important nodes point to them. A high-PageRank service is one that many other significant services connect to.
  • Betweenness centrality measures how often a node sits on the shortest path between other nodes. High-betweenness nodes are "stepping stones" attackers must traverse.
  • In-degree counts direct incoming connections. Nodes with many inbound edges are natural targets because many hosts connect to them.
  • Run the cell. You should see a ranked list of your top 10 nodes by PageRank score, showing which service endpoints are structurally most important in your knowledge graph.

Seeing a PowerIterationFailedConvergence error?

This means PageRank did not converge within the default 100 iterations. Try increasing the iteration limit by passing max_iter=200 to nx.pagerank(G, max_iter=200).

help me fix a PowerIterationFailedConvergence error in NetworkX pagerank

Implement threat scoring

Now you need to combine the centrality metrics with what the GNN learned about each node. The GNN's hidden layer stores an embedding vector for each node that encodes its neighborhood context.

  • Extract the GNN's learned node embeddings and identify known attack targets by adding this cell:
# Extract GNN node embeddings from the first layer (before classification)
gnn_model.eval()
with torch.no_grad():
    embeddings = F.relu(gnn_model.conv1(data.x, data.edge_index))

# Identify known attack target nodes from test set
test_indices = data.test_mask.nonzero(as_tuple=True)[0]
attack_mask = data.y[test_indices] > 0
known_attack_nodes = test_indices[attack_mask].tolist()

# Store node list for name lookups
node_list = list(G.nodes())

print(f"Embeddings shape: {embeddings.shape}")
print(f"Known attack target nodes: {len(known_attack_nodes)}")

What does this code do?

  • gnn_model.conv1 is the first GraphSAGE layer. Running the forward pass through only this layer (without the classification head) gives you a 64-dimensional embedding for each node.
  • known_attack_nodes collects every test-set node that was labeled as an attack (label > 0). These are the nodes your threat scorer will compare against.
  • Run the cell. You should see the embeddings shape printed (number of nodes by 64 dimensions) and the count of known attack target nodes.
  • Define the threat prediction function and test it by adding this cell:
def predict_threat_score(node_id):
    """Score a node's threat level using centrality + embedding similarity."""
    # Get the target node's embedding
    node_embedding = embeddings[node_id].unsqueeze(0)

    # Compute cosine similarity to known attack target embeddings
    attack_embeddings = embeddings[known_attack_nodes]
    similarities = F.cosine_similarity(node_embedding, attack_embeddings)
    avg_similarity = similarities.mean().item()

    # Look up centrality metrics for this node
    node_name = node_list[node_id]
    pr_score = pagerank_scores.get(node_name, 0)
    bc_score = betweenness_scores.get(node_name, 0)
    deg_score = indegree_scores.get(node_name, 0)

    # Weighted combination of similarity and centrality
    max_deg = max(indegree_scores.values()) if max(indegree_scores.values()) > 0 else 1
    threat_score = (0.4 * avg_similarity) + (0.3 * pr_score * 100) + (0.2 * bc_score) + (0.1 * (deg_score / max_deg))
    return threat_score

# Quick test on first non-attack node
test_idx = [i for i in range(len(node_list)) if i not in known_attack_nodes][0]
print(f"Test threat score for '{node_list[test_idx]}': {predict_threat_score(test_idx):.4f}")

How does the scoring formula work?

The function blends four signals into a single threat score. Cosine similarity (weight 0.4) measures how close a node's learned representation is to known attack targets in the GNN's embedding space.

PageRank (weight 0.3) rewards high-value targets. Betweenness centrality (weight 0.2) flags stepping-stone nodes. In-degree (weight 0.1) captures raw exposure.

  • Run the cell. You should see a test threat score printed for one node, confirming the function works end-to-end.

Getting an IndexError on embeddings?

Make sure you ran the previous cell that computes embeddings and known_attack_nodes. The function references these variables from the notebook's global scope.

help me debug an IndexError in my predict_threat_score function

Validate predictions against test data

A threat score is only useful if it actually predicts real attacks. You need to verify that nodes your scorer ranks highly were indeed attacked in the held-out test set.

  • Score all non-attack nodes and generate a ranked threat list by adding this cell:
# Score all non-attack nodes
threat_scores = {}
for idx in range(len(node_list)):
    if idx not in known_attack_nodes:
        threat_scores[idx] = predict_threat_score(idx)

# Rank by threat score (highest first)
ranked_threats = sorted(threat_scores.items(), key=lambda x: x[1], reverse=True)

# Print ranked list of services with threat scores
print("Top 10 predicted threat targets:")
print("-" * 40)
for idx, score in ranked_threats[:10]:
    print(f"  {node_list[idx]}: threat score = {score:.4f}")
  • Run the cell. You should see a ranked list of 10 services with their computed threat scores, ordered from highest to lowest.

Before you run the validation, think about this: do you expect the top-ranked services to actually correspond to real attack targets in the test data?

  • Validate the predictions against actual attack targets by adding this final cell:
# Validate: check if top-ranked nodes were actual attack targets
top_k = 20
top_predicted = [idx for idx, _ in ranked_threats[:top_k]]

# Check which predicted targets have attack labels
hits = sum(1 for idx in top_predicted if data.y[idx].item() > 0)
precision_at_k = hits / top_k

print(f"Validation Results:")
print(f"  Top-{top_k} predicted targets evaluated")
print(f"  Actual attack targets found: {hits}")
print(f"  Precision@{top_k}: {precision_at_k:.2f}")
print(f"\nGraph-based threat prediction identifies services")
print(f"likely to be targeted based on structural patterns.")

What does Precision@k mean here?

Precision@20 tells you what fraction of your top 20 predicted targets were actually attacked in the test set. A precision of 0.5 or higher means the graph-based reasoning identifies real attack targets at a rate well above random chance.

This demonstrates that the combination of graph centrality and GNN embeddings captures structural attack patterns that can forecast future targets.

  • Run the cell. You should see the validation results showing how many of your top-20 predicted targets were real attack targets, with a Precision@20 score above random baseline.

Precision@20 is 0.0?

If all predictions miss, check that your data.y labels are correctly aligned with node_list. The label at index i in the tensor must correspond to node i in the graph's node ordering.

Also verify that you are filtering known_attack_nodes correctly. The scorer should only evaluate nodes NOT already in the known attack set.

help me debug why my threat prediction precision is zero

✔️ Awesome, I've got everything!

Great. Make sure all cells are saved in your notebook.

ⓧ I'd like to double check the full code

Scroll up through your notebook and confirm the following cells are present in order:

  • Cell computing pagerank_scores, betweenness_scores, and indegree_scores with the top-10 PageRank printout.
  • Cell extracting embeddings from gnn_model.conv1 and building known_attack_nodes.
  • Cell defining predict_threat_score(node_id) with a test print at the end.
  • Cell scoring all nodes and printing the ranked threat list.
  • Cell validating predictions against test labels and printing Precision@20.

Secret mission

Enrich with MITRE ATT&CK Mappings

Your intrusion detector flags anomalies, but a real SOC analyst needs more than a label. They need to know which tactic an attacker is using, what technique maps to the alert, and which other services face the same threat. In this secret mission, you'll add a MITRE ATT&CK ontology layer to your knowledge graph and build graph traversal queries that turn raw detections into actionable threat intelligence.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • intrusion_detection.ipynb notebook file.
  • NSL-KDD dataset CSV files.
  • Python packages: torch, torch-geometric, pyg-lib, torch-scatter, torch-sparse, networkx, scikit-learn, pandas, matplotlib, seaborn.

Keep everything running

No action needed. Choose this if you plan to continue experimenting with your knowledge graph or want to extend the model further.

  • Your notebook and dataset remain in place, ready for further development.
  • All installed Python packages stay available for future graph-based ML projects.

Pause - I'll come back to this later

Shut down running processes to free up memory. Your files and packages stay installed so you can pick up where you left off.

  • Close the Jupyter Notebook kernel by clicking Kernel then Shut Down Kernel in VS Code.
  • Close the intrusion_detection.ipynb tab in VS Code.
  • Your dataset and packages remain installed for when you return.

Delete - I don't want to use this again

Remove all project files and uninstall the packages to start fresh.

Delete the project files by running the appropriate command for your operating system.

macOS

Delete the notebook and dataset files by running:

rm intrusion_detection.ipynb
rm *.csv

Windows

Delete the notebook and dataset files by running:

del intrusion_detection.ipynb
del *.csv

Uninstall the Python packages by running:

pip uninstall torch torch-geometric pyg-lib torch-scatter torch-sparse networkx scikit-learn pandas matplotlib seaborn -y

Nice Work!

Nice Work!

You just built a knowledge graph-powered intrusion detection system from scratch. From raw network traffic data to graph neural network predictions, you demonstrated how knowledge engineering transforms a classification problem into a structural reasoning problem.

You've learned how to:

  • Construct a knowledge graph from raw network traffic data, designing an ontology with hosts, services, and protocols as entities connected by typed relationships.
  • Train a GraphSAGE GNN that leverages graph topology to improve detection of rare attack types (R2L, U2R) where flat classifiers like Random Forest fail.
  • Implement threat prediction using graph centrality metrics and learned node embeddings to forecast which services will be targeted next.
  • (Secret Mission) Enriched the knowledge graph with MITRE ATT&CK mappings to provide tactical context for detected intrusions, demonstrating knowledge engineering principles applicable to real SOC operations.

Ready to quiz yourself?