# FedDGCN: A Scalable Federated Learning Framework for Traffic Flow Prediction

This is the official repository of **FedDGCN:  A Scalable Federated Learning Framework for Traffic Flow Prediction.**  

![Overview](./figures/overview.jpg)

**FedDGCN** extends [FederatedScope](https://github.com/alibaba/FederatedScope) to support federated traffic flow prediction.  

> **Note:** This is an early version of **FedDGCN**. The full version will be released after testing is completed.

---

## Table of Contents  
- [1. Environment Setup](#1-environment-setup)  
  - [Step 1: Create a Conda Environment](#step-1-create-a-conda-environment)  
  - [Step 2: Install PyTorch](#step-2-install-pytorch)  
  - [Step 3: Install FederatedScope](#step-3-install-federatedscope)  
- [2. Run the Code](#2-run-the-code)  
  - [Step 1: Prepare the Datasets](#step-1-prepare-the-datasets)  
  - [Step 2: Configure the Settings](#step-2-configure-the-settings)  
  - [Step 3: Run the Experiments](#step-3-run-the-experiments)  
- [3. Visualize Results](#3-visualize-results)  
- [4. Citation](#4-citation)  
- [5. Acknowledgements](#5-acknowledgements)  

---

## 1. Environment Setup

### Step 1: Create a Conda Environment  
We recommend using a **Conda** virtual environment. This project supports **Python 3.9** (recommended) and **Python 3.10**.  

> **Warning:** Python 3.11 and later versions are not compatible.  

```bash
conda create -n FedDGCN python=3.9
conda activate FedDGCN
```

### Step 2: Install PyTorch  
Download the appropriate version of [PyTorch](https://pytorch.org/get-started/locally/) based on your device.  

This project has been tested with **Torch 2.4.0 (recommended)** and **Torch 2.0.0** with **CUDA 12**. Compatibility with other versions is not guaranteed.  

### Step 3: Install FederatedScope  
Clone this repository and install it:  

```bash
git clone https://github.com/your-repo/FS-TFP.git
cd FS-TFP
pip install -e .
```

Additionally, install the required packages to avoid warnings:  

```bash
pip install torch_geometric community rdkit
```

---

## 2. Run the Code  

### Step 1: Prepare the Datasets  
Download the PeMS datasets from the **[STSGCN repository](https://github.com/Davidham3/STSGCN)**.  
After downloading, extract the datasets and place them in the `./data/trafficflow` directory.  

The directory structure should be as follows:  

```
FS-TFP/data/trafficflow
├─PeMS03
├─PeMS04
├─PeMS07
└─PeMS08
```

### Step 2: Configure the Settings  
Run scripts for the four datasets are located in the `./scripts/trafficflow_exp_scripts/` directory.  

Each dataset has a YAML configuration file: `{D3, D4, D7, D8}.yaml`.  

You can customize the parameters or use the presets. Key configurable parameters include:  

```yaml
# Line 3: GPU device to use (for multi-GPU machines)
device: 0  

# Line 8: Total number of training rounds
total_round_num: <number_of_rounds>

# Line 9: Number of clients
client_num: <number_of_clients>

# Line 65: Training loss function
# Options: L1Loss, RMSE, MAPE
criterion:
  type: <loss_function>
```

> **Warning:** Processing the **PeMSD7** dataset may require more than **32GB RAM**.  
> If your system lacks sufficient RAM, increase the size of the swap partition.  

### Step 3: Run the Experiments  

Use the following commands to run **FedDGCN**:  

```bash
# Run experiments on different datasets
python federatedscope/main.py --cfg scripts/trafficflow_exp_scripts/D3.yaml  # PeMSD3
python federatedscope/main.py --cfg scripts/trafficflow_exp_scripts/D4.yaml  # PeMSD4
python federatedscope/main.py --cfg scripts/trafficflow_exp_scripts/D7.yaml  # PeMSD7
python federatedscope/main.py --cfg scripts/trafficflow_exp_scripts/D8.yaml  # PeMSD8
```

If you see output similar to the image below, congratulations! You have successfully run the experiment:  

![Experiment Output](./figures/exp.png)

---

## 3. Visualize Results  

The experiment logs will be saved in the **`exp`** folder.  

We provide a script, **`global.py`**, in the same folder. Replace the old logs with the new logs from the experiments and run the script to visualize the results:  

```bash
python exp/global.py
```

This will generate a **baseline.jpg** file to visualize the logs.  

To install the required package for visualization:  

```bash
pip install matplotlib
```

---

## 4. Citation  

TBD

---

## 5. Acknowledgements  

Special thanks to the authors of [FederatedScope](https://github.com/alibaba/FederatedScope), upon which this project is built.  

We also express our gratitude to the following projects: [AGCRN](https://github.com/LeiBAI/AGCRN), [STG-NCDE](https://github.com/jeongwhanchoi/STG-NCDE), [RGDAN](https://github.com/wengwenchao123/RGDAN).




# How to improve our framework?

We welcome the community to help expand and improve our framework! This guide outlines how to customize configurations, models, datasets, trainers, and loss functions.

---


## 📋 How to Add Configurations

1. **Create a YAML file**  
   Use the `./scripts` folder as a reference. We recommend copying an existing configuration and modifying it.

2. **Update Core Configurations**  
   Go to `./federatedscope/core/configs/`, locate `cfg_model`, and add your configurations with default values. For example:
   ```python
   cfg.model.num_nodes = 0
   cfg.model.rnn_units = 64
   cfg.model.dropout = 0.1
   ```

3. **Add Nested Parameters**  
   If needed, create nested parameters using `CN()`:
   ```python
   cfg.model.next = CN()
   cfg.model.next.default = 1
   ```

4. **Sync Parameters Across Configs**  
   Ensure the parameters in `cfg_model` align with those in other configs (e.g., `cfg_trafficflow`) to avoid compatibility issues across systems (Windows, Linux, etc.).

5. **Customize YAML**  
   After adding parameters to config files (e.g., `cfg_data`, `cfg_training`), customize them in the corresponding YAML file.

---

## 🛠️ How to Add a Model

1. **Create Your Model**  
   Save your model in the `federatedscope/trafficflow/model/` folder (or another location of your choice).  
   Add the model name to `model:type` in the configuration YAML file.

2. **Register the Model**  
   In `federatedscope/core/auxiliaries/model_builder.py`, add logic for your model (around line 214):
   ```python
   elif model_config.type.lower() in ['your_model']:
       from federatedscope.trafficflow.model.your_model import YourModel
       model = YourModel(model_config)
   ```

---

## 📊 How to Add a Dataset

1. **Create a DataLoader**  
   Implement a function in `federatedscope/trafficflow/dataloader/` to generate data for clients. The function should return a list of dictionaries (one per client) with `['train']`, `['val']`, and `['test']` datasets. Each dataset should be a `torch.utils.data.TensorDataset` containing `x` and `label` tensors.

2. **Register the DataLoader**  
   Update `federatedscope/core/data/utils.py` (around line 108) to import your dataloader:
   ```python
   elif config.data.type.lower() in ['trafficflow']:
       from federatedscope.trafficflow.dataloader.traffic_dataloader import load_traffic_data
       dataset, modified_config = load_traffic_data(config, client_cfgs)
   ```

---

## 🎓 How to Customize Your Trainer

1. **Create a Custom Trainer**  
   Implement your trainer in `federatedscope/trafficflow/trainer/`. Inherit from an existing trainer, e.g.:
   ```python
   from federatedscope.core.trainers.torch_trainer import GeneralTorchTrainer as Trainer
   
   class TrafficflowTrainer(Trainer):
       # Overwrite methods or hooks as needed
   ```

2. **Reference Examples**  
   Review existing trainers for additional guidance.

---

## 🔧 How to Add a Custom Loss Function

1. **Implement the Loss Function**  
   Create a file in `federatedscope/contrib/loss/` and write your custom loss function.

2. **Register the Loss Function**  
   Register your loss function using `register_criterion`:
   
   ```python
   from federatedscope.register import register_criterion
   
   register_criterion('RMSE', call_my_criterion)
   register_criterion('MAPE', call_my_criterion)
   ```
   
3. **Update Configuration**  
   Set the loss function in the configuration YAML file using `criterion:type`.

---

By following these steps, you can extend and customize the framework to meet your needs. Happy coding! 🎉