Federated Learning with Network-Aware GNN

A federated learning system for cybersecurity attack prediction using Graph Neural Networks.

---

🎯 Overview

Train a GNN model across multiple clients without sharing raw data. Predicts the next attack technique in a cybersecurity attack chain.

Tech Stack: PyTorch | PyTorch Geometric | Flower FL

---

📁 Project Structure

.
├── config.py                   # All configuration settings
├── requirements.txt            # Python dependencies
├── docker-compose.yml          # Docker setup
│
├── models/                     # Pre-trained model files
│   ├── best_network_aware_model.pt
│   └── network_aware_mappings.pkl
│
├── data/                       # Client datasets (auto-generated)
│   ├── client1/
│   ├── client2/
│   └── client3/
│
├── client/                     # Client implementation
│   ├── client.py
│   └── model_loader.py
│
├── server/                     # Server implementation
│   └── server.py
│
├── scripts/
│   └── split_gnn_data.py      # Split data for clients
│
├── start_server.py            # Launch server
└── start_client.py            # Launch client

---

🚀 Quick Start

Local Setup

# 1. Create environment
python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate

# 2. Install dependencies
pip install -r requirements.txt

# 3. Prepare data
python scripts/split_gnn_data.py

# 4. Start training (4 terminals)
python start_server.py           # Terminal 1
python start_client.py 1         # Terminal 2
python start_client.py 2         # Terminal 3
python start_client.py 3         # Terminal 4

Docker Setup

# 1. Prepare data (if not done)
python scripts/split_gnn_data.py

# 2. Start all containers
docker-compose up --build

# 3. Monitor logs
docker-compose logs -f

# 4. Stop
docker-compose down

---

⚙️ Configuration

Edit config.py:

# Federated Learning
NUM_CLIENTS = 3                 # Number of clients
NUM_ROUNDS = 10                 # Training rounds
MIN_CLIENTS = 2                 # Min clients per round

# Training
BATCH_SIZE = 32
LOCAL_EPOCHS = 3
LEARNING_RATE = 0.001
DEVICE = "cpu"                  # or "cuda" for GPU

---

📊 Metrics

The system tracks:

• Loss: Cross-entropy loss
• Top-1 Accuracy: Exact match
• Top-3 Accuracy: Target in top 3
• Top-5 Accuracy: Target in top 5

---

🔧 Troubleshooting

Problem	Solution
Model not found	Copy from repogit/Predection/models/models/ to ./models/
Data not found	Run python scripts/split_gnn_data.py
Import errors	pip install torch torch-geometric flwr
CUDA out of memory	Set DEVICE="cpu" or reduce BATCH_SIZE
Docker fails	docker-compose down -v && docker-compose up --build

---

📈 Performance Tips

Better Accuracy:

• Increase NUM_ROUNDS (20-50)
• Increase LOCAL_EPOCHS (5-10)
• Use more training data

Faster Training:

• Use GPU (DEVICE="cuda")
• Increase BATCH_SIZE
• Reduce LOCAL_EPOCHS

---

📚 Additional Documentation

• SYSTEM_OVERVIEW.md - Architecture details
• DEPLOYMENT.md - Deployment guides
• FL_GNN_README.md - Technical details

---

📄 License

MIT License

---

Built with 🧠 PyTorch | 🌸 Flower | 🐳 Docker