Federated Link Prediction Example

Federated Link Prediction with STFL on the Link Prediction dataset.

(Time estimate: 3 minutes)

Load libraries

import os

import attridict

from fedgraph.federated_methods import run_fedgraph

Specify the Link Prediction configuration

BASE_DIR = os.path.dirname(os.path.abspath("."))
DATASET_PATH = os.path.join(
    BASE_DIR, "data", "LPDataset"
)  # Could be modified based on the user needs
config = {
    "fedgraph_task": "LP",
    # method = ["STFL", "StaticGNN", "4D-FED-GNN+", "FedLink"]
    "method": "STFL",
    # Dataset configuration
    # country_codes = ['US', 'BR', 'ID', 'TR', 'JP']
    "country_codes": ["JP"],
    "dataset_path": DATASET_PATH,
    # Setup configuration
    "device": "cpu",
    "use_buffer": False,
    "buffer_size": 300000,
    "online_learning": False,
    "seed": 10,
    # Model parameters
    "global_rounds": 8,
    "local_steps": 3,
    "hidden_channels": 64,
    # Output configuration
    "record_results": False,
    # System configuration
    "gpu": False,
    "num_cpus_per_trainer": 1,
    "num_gpus_per_trainer": 0,
    "use_cluster": False,  # whether use kubernetes for scalability or not
    "distribution_type": "average",  # the node number distribution among clients
    "batch_size": -1,  # -1 is full batch
}

Run fedgraph method

config = attridict(config)
run_fedgraph(config)

/home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset not exists, creating directory
Downloading traveled_users from https://drive.google.com/uc?id=1RUsyGrsz4hmY3OA3b-oqyh5yqlks02-p...
Downloading...
From: https://drive.google.com/uc?id=1RUsyGrsz4hmY3OA3b-oqyh5yqlks02-p
To: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/traveled_users.txt

  0%|          | 0.00/552k [00:00<?, ?B/s]
 95%|█████████▍| 524k/552k [00:00<00:00, 4.74MB/s]
100%|██████████| 552k/552k [00:00<00:00, 4.96MB/s]
Downloaded traveled_users to /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/traveled_users.txt
Downloading data_global from https://drive.google.com/uc?id=1CnBlVXqCbfjSswagTci5D7nAqO7laU_J...
Downloading...
From (original): https://drive.google.com/uc?id=1CnBlVXqCbfjSswagTci5D7nAqO7laU_J
From (redirected): https://drive.google.com/uc?id=1CnBlVXqCbfjSswagTci5D7nAqO7laU_J&confirm=t&uuid=4b94c4a5-d8ae-4b21-a4df-11f309325fda
To: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_global.txt

  0%|          | 0.00/278M [00:00<?, ?B/s]
  0%|          | 524k/278M [00:00<00:58, 4.71MB/s]
  4%|▎         | 9.96M/278M [00:00<00:04, 55.0MB/s]
  8%|▊         | 22.5M/278M [00:00<00:02, 86.1MB/s]
 11%|█▏        | 31.5M/278M [00:00<00:03, 77.4MB/s]
 18%|█▊        | 48.8M/278M [00:00<00:02, 107MB/s]
 24%|██▎       | 65.5M/278M [00:00<00:01, 125MB/s]
 28%|██▊       | 78.6M/278M [00:00<00:02, 97.5MB/s]
 33%|███▎      | 90.7M/278M [00:00<00:02, 93.6MB/s]
 38%|███▊      | 105M/278M [00:01<00:01, 102MB/s]
 44%|████▍     | 123M/278M [00:01<00:01, 120MB/s]
 50%|████▉     | 138M/278M [00:01<00:01, 127MB/s]
 56%|█████▌    | 155M/278M [00:01<00:00, 138MB/s]
 62%|██████▏   | 173M/278M [00:01<00:00, 149MB/s]
 68%|██████▊   | 190M/278M [00:01<00:00, 153MB/s]
 74%|███████▍  | 206M/278M [00:01<00:00, 152MB/s]
 80%|███████▉  | 221M/278M [00:01<00:00, 144MB/s]
 86%|████████▌ | 239M/278M [00:01<00:00, 152MB/s]
 92%|█████████▏| 256M/278M [00:02<00:00, 156MB/s]
 98%|█████████▊| 272M/278M [00:02<00:00, 148MB/s]
100%|██████████| 278M/278M [00:02<00:00, 126MB/s]
Downloaded data_global to /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_global.txt
Downloading data_JP from https://drive.google.com/uc?id=1IPBW4dRYk52x8TahfBqFOh3GdxoYafJ2...
Downloading...
From: https://drive.google.com/uc?id=1IPBW4dRYk52x8TahfBqFOh3GdxoYafJ2
To: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_JP.txt

  0%|          | 0.00/28.7M [00:00<?, ?B/s]
  2%|▏         | 524k/28.7M [00:00<00:06, 4.67MB/s]
 33%|███▎      | 9.44M/28.7M [00:00<00:00, 51.6MB/s]
 53%|█████▎    | 15.2M/28.7M [00:00<00:00, 39.4MB/s]
100%|██████████| 28.7M/28.7M [00:00<00:00, 58.0MB/s]
Downloaded data_JP to /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_JP.txt
gpu not detected
start training
global rounds: 0
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) checking code and file path: JP,/home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset
(Trainer pid=1052) printing in getdata, path: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset
(Trainer pid=1052) Loading data in /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_JP.txt
(Trainer pid=1052) Device: 'cpu'
(Trainer pid=1052) loading train_data and test_data
(Trainer pid=1052) client 0 local steps 0 loss 0.7616 train time 6.2408
(Trainer pid=1052) client 0 local steps 1 loss 0.6568 train time 6.0289
clientId: 0 current_loss: 0.5707166194915771 train_finish_times: [6.240787744522095, 6.028871536254883, 5.897978067398071]
(Trainer pid=1052) client 0 local steps 2 loss 0.5707 train time 5.8980
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.7273755669593811 hit rate: 0.8597114086151123 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.7273755669593811 hit rate: 0.8597114086151123 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.7273755669593811 hit rate: 0.8597114086151123
global rounds: 1
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) Test AUC: 0.7274
(Trainer pid=1052) Test Hit Rate at 2: 0.8597
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1052) client 0 local steps 0 loss 0.4959 train time 5.9981
(Trainer pid=1052) client 0 local steps 1 loss 0.4307 train time 5.8977
clientId: 0 current_loss: 0.37483564019203186 train_finish_times: [5.9981255531311035, 5.897661447525024, 5.876627206802368]
(Trainer pid=1052) client 0 local steps 2 loss 0.3748 train time 5.8766
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.8073558807373047 hit rate: 0.9038324952125549 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.8073558807373047 hit rate: 0.9038324952125549 traveled user hit rate: 1.0

(Trainer pid=1052) Test AUC: 0.8074
(Trainer pid=1052) Test Hit Rate at 2: 0.9038
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
Predict Day 20 average auc score: 0.8073558807373047 hit rate: 0.9038324952125549
global rounds: 2
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) client 0 local steps 0 loss 0.3276 train time 5.9780
(Trainer pid=1052) client 0 local steps 1 loss 0.2885 train time 5.9535
clientId: 0 current_loss: 0.2570096552371979 train_finish_times: [5.978019952774048, 5.953490734100342, 5.847478628158569]
(Trainer pid=1052) client 0 local steps 2 loss 0.2570 train time 5.8475
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.8171350955963135 hit rate: 0.910338282585144 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.8171350955963135 hit rate: 0.910338282585144 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8171350955963135 hit rate: 0.910338282585144
(Trainer pid=1052) Test AUC: 0.8171
(Trainer pid=1052) Test Hit Rate at 2: 0.9103
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
global rounds: 3
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) client 0 local steps 0 loss 0.2330 train time 6.0237
(Trainer pid=1052) client 0 local steps 1 loss 0.2154 train time 5.8747
clientId: 0 current_loss: 0.20293746888637543 train_finish_times: [6.023711204528809, 5.874679088592529, 6.041338682174683]
(Trainer pid=1052) client 0 local steps 2 loss 0.2029 train time 6.0413
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.8197392225265503 hit rate: 0.9108114242553711 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.8197392225265503 hit rate: 0.9108114242553711 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8197392225265503 hit rate: 0.9108114242553711
global rounds: 4
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) Test AUC: 0.8197
(Trainer pid=1052) Test Hit Rate at 2: 0.9108
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1052) client 0 local steps 0 loss 0.1940 train time 6.2308
(Trainer pid=1052) client 0 local steps 1 loss 0.1874 train time 6.1495
clientId: 0 current_loss: 0.18198850750923157 train_finish_times: [6.230805158615112, 6.149538516998291, 6.255640029907227]
(Trainer pid=1052) client 0 local steps 2 loss 0.1820 train time 6.2556
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.8215453624725342 hit rate: 0.9131771922111511 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.8215453624725342 hit rate: 0.9131771922111511 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8215453624725342 hit rate: 0.9131771922111511
global rounds: 5
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) Test AUC: 0.8215
(Trainer pid=1052) Test Hit Rate at 2: 0.9132
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1052) client 0 local steps 0 loss 0.1772 train time 6.2600
(Trainer pid=1052) client 0 local steps 1 loss 0.1727 train time 6.1995
clientId: 0 current_loss: 0.16825395822525024 train_finish_times: [6.259976863861084, 6.199476480484009, 6.228652477264404]
(Trainer pid=1052) client 0 local steps 2 loss 0.1683 train time 6.2287
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.823245644569397 hit rate: 0.9138869047164917 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.823245644569397 hit rate: 0.9138869047164917 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.823245644569397 hit rate: 0.9138869047164917
global rounds: 6
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) Test AUC: 0.8232
(Trainer pid=1052) Test Hit Rate at 2: 0.9139
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1052) client 0 local steps 0 loss 0.1639 train time 6.2523
(Trainer pid=1052) client 0 local steps 1 loss 0.1597 train time 6.2454
clientId: 0 current_loss: 0.15555423498153687 train_finish_times: [6.252341270446777, 6.245421409606934, 6.077695846557617]
(Trainer pid=1052) client 0 local steps 2 loss 0.1556 train time 6.0777
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.8248248100280762 hit rate: 0.9144783616065979 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.8248248100280762 hit rate: 0.9144783616065979 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8248248100280762 hit rate: 0.9144783616065979
global rounds: 7
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1052) Test AUC: 0.8248
(Trainer pid=1052) Test Hit Rate at 2: 0.9145
(Trainer pid=1052) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1052) client 0 local steps 0 loss 0.1516 train time 6.3092
(Trainer pid=1052) client 0 local steps 1 loss 0.1479 train time 6.1430
clientId: 0 current_loss: 0.14426136016845703 train_finish_times: [6.309248685836792, 6.142970323562622, 6.080567836761475]
(Trainer pid=1052) client 0 local steps 2 loss 0.1443 train time 6.0806
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) auc score: 0.8262328505516052 hit rate: 0.9150698184967041 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, 32df7081231e7a5906b4255201000000) final auc score: 0.8262328505516052 hit rate: 0.9150698184967041 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8262328505516052 hit rate: 0.9150698184967041
training is not complete
The whole process has ended

Total running time of the script: (3 minutes 39.320 seconds)

Gallery generated by Sphinx-Gallery