Note
Go to the end to download the full example code.
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)
Initialization start time recorded.
/home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/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/stable/data/LPDataset/traveled_users.txt
0%| | 0.00/552k [00:00<?, ?B/s]
100%|██████████| 552k/552k [00:00<00:00, 12.8MB/s]
Downloaded traveled_users to /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/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=85e21281-f0a8-4029-8d05-3685287d720c
To: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/data/LPDataset/data_global.txt
0%| | 0.00/278M [00:00<?, ?B/s]
2%|▏ | 6.82M/278M [00:00<00:04, 66.5MB/s]
5%|▌ | 15.2M/278M [00:00<00:03, 76.1MB/s]
9%|▉ | 24.6M/278M [00:00<00:03, 82.6MB/s]
12%|█▏ | 34.6M/278M [00:00<00:02, 87.7MB/s]
16%|█▌ | 43.5M/278M [00:00<00:02, 85.8MB/s]
19%|█▉ | 52.4M/278M [00:00<00:02, 86.1MB/s]
22%|██▏ | 61.9M/278M [00:00<00:02, 88.5MB/s]
25%|██▌ | 70.8M/278M [00:00<00:02, 87.0MB/s]
29%|██▊ | 79.7M/278M [00:00<00:02, 87.0MB/s]
32%|███▏ | 88.6M/278M [00:01<00:02, 82.6MB/s]
35%|███▌ | 98.0M/278M [00:01<00:02, 84.3MB/s]
38%|███▊ | 107M/278M [00:01<00:02, 79.2MB/s]
41%|████▏ | 115M/278M [00:01<00:02, 79.3MB/s]
45%|████▍ | 124M/278M [00:01<00:01, 79.1MB/s]
48%|████▊ | 133M/278M [00:01<00:01, 82.3MB/s]
51%|█████▏ | 143M/278M [00:01<00:01, 83.4MB/s]
54%|█████▍ | 151M/278M [00:01<00:01, 83.1MB/s]
57%|█████▋ | 159M/278M [00:01<00:01, 79.7MB/s]
60%|██████ | 168M/278M [00:02<00:01, 77.4MB/s]
63%|██████▎ | 176M/278M [00:02<00:01, 77.3MB/s]
66%|██████▌ | 184M/278M [00:02<00:01, 73.3MB/s]
69%|██████▉ | 192M/278M [00:02<00:01, 74.6MB/s]
72%|███████▏ | 201M/278M [00:02<00:00, 78.6MB/s]
75%|███████▌ | 209M/278M [00:02<00:00, 77.1MB/s]
78%|███████▊ | 218M/278M [00:02<00:00, 78.9MB/s]
81%|████████▏ | 226M/278M [00:02<00:00, 80.9MB/s]
85%|████████▍ | 235M/278M [00:02<00:00, 81.8MB/s]
88%|████████▊ | 244M/278M [00:03<00:00, 82.2MB/s]
91%|█████████ | 252M/278M [00:03<00:00, 81.8MB/s]
94%|█████████▎| 261M/278M [00:03<00:00, 82.3MB/s]
97%|█████████▋| 269M/278M [00:03<00:00, 80.6MB/s]
100%|█████████▉| 277M/278M [00:03<00:00, 81.4MB/s]
100%|██████████| 278M/278M [00:03<00:00, 81.2MB/s]
Downloaded data_global to /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/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/stable/data/LPDataset/data_JP.txt
0%| | 0.00/28.7M [00:00<?, ?B/s]
18%|█▊ | 5.24M/28.7M [00:00<00:00, 51.5MB/s]
40%|████ | 11.5M/28.7M [00:00<00:00, 57.3MB/s]
69%|██████▉ | 19.9M/28.7M [00:00<00:00, 62.8MB/s]
99%|█████████▊| 28.3M/28.7M [00:00<00:00, 69.1MB/s]
100%|██████████| 28.7M/28.7M [00:00<00:00, 65.9MB/s]
Downloaded data_JP to /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/data/LPDataset/data_JP.txt
gpu not detected
[Debug] Server running on IP: 172.17.0.2
//Log init_time: 22009.697 ms //end
Pretrain start time recorded.
//pretrain_time: 749.922 ms//end
Train start time recorded.
start training
global rounds: 0
Training in LP_train_global_round, number of clients: 1
(pid=1490) FedGraph Differential Privacy module loaded (v1.0.0)
(pid=1490) ✓ Differential Privacy support loaded
(Trainer pid=1490) checking code and file path: JP,/home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/data/LPDataset
(Trainer pid=1490) printing in getdata, path: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/data/LPDataset
(Trainer pid=1490) Loading data in /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/stable/data/LPDataset/data_JP.txt
(Trainer pid=1490) Device: 'cpu'
(Trainer pid=1490) [Debug] Trainer running on node IP: 172.17.0.2
(Trainer pid=1490) loading train_data and test_data
(Trainer pid=1490) client 0 local steps 0 loss 0.7619 train time 7.0120
(Trainer pid=1490) client 0 local steps 1 loss 0.6569 train time 6.7213
clientId: 0 current_loss: 0.5707483887672424 train_finish_times: [7.012034177780151, 6.721317768096924, 6.696970701217651]
(Trainer pid=1490) client 0 local steps 2 loss 0.5707 train time 6.6970
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.7258930206298828 hit rate: 0.8543884754180908 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.7258930206298828 hit rate: 0.8543884754180908 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.7258930206298828 hit rate: 0.8543884754180908
global rounds: 1
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) Test AUC: 0.7259
(Trainer pid=1490) Test Hit Rate at 2: 0.8544
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1490) client 0 local steps 0 loss 0.4958 train time 6.7320
(Trainer pid=1490) client 0 local steps 1 loss 0.4305 train time 6.6807
clientId: 0 current_loss: 0.3745068907737732 train_finish_times: [6.731971979141235, 6.680720567703247, 6.726215362548828]
(Trainer pid=1490) client 0 local steps 2 loss 0.3745 train time 6.7262
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8064392805099487 hit rate: 0.9021764993667603 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8064392805099487 hit rate: 0.9021764993667603 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8064392805099487 hit rate: 0.9021764993667603
(Trainer pid=1490) Test AUC: 0.8064
(Trainer pid=1490) Test Hit Rate at 2: 0.9022
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
global rounds: 2
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) client 0 local steps 0 loss 0.3272 train time 6.6628
(Trainer pid=1490) client 0 local steps 1 loss 0.2879 train time 6.7259
clientId: 0 current_loss: 0.2563645839691162 train_finish_times: [6.6627771854400635, 6.725903034210205, 6.727693319320679]
(Trainer pid=1490) client 0 local steps 2 loss 0.2564 train time 6.7277
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8165777325630188 hit rate: 0.9115211963653564 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8165777325630188 hit rate: 0.9115211963653564 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8165777325630188 hit rate: 0.9115211963653564
global rounds: 3
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) Test AUC: 0.8166
(Trainer pid=1490) Test Hit Rate at 2: 0.9115
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1490) client 0 local steps 0 loss 0.2322 train time 6.6688
(Trainer pid=1490) client 0 local steps 1 loss 0.2146 train time 6.6935
clientId: 0 current_loss: 0.20218336582183838 train_finish_times: [6.66879677772522, 6.693462133407593, 6.660843372344971]
(Trainer pid=1490) client 0 local steps 2 loss 0.2022 train time 6.6608
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8193737864494324 hit rate: 0.9104565978050232 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8193737864494324 hit rate: 0.9104565978050232 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8193737864494324 hit rate: 0.9104565978050232
global rounds: 4
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) Test AUC: 0.8194
(Trainer pid=1490) Test Hit Rate at 2: 0.9105
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1490) client 0 local steps 0 loss 0.1933 train time 6.6978
(Trainer pid=1490) client 0 local steps 1 loss 0.1867 train time 6.6873
clientId: 0 current_loss: 0.18133950233459473 train_finish_times: [6.6977763175964355, 6.687315940856934, 6.679613828659058]
(Trainer pid=1490) client 0 local steps 2 loss 0.1813 train time 6.6796
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8213432431221008 hit rate: 0.9125857353210449 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8213432431221008 hit rate: 0.9125857353210449 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8213432431221008 hit rate: 0.9125857353210449
global rounds: 5
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) Test AUC: 0.8213
(Trainer pid=1490) Test Hit Rate at 2: 0.9126
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1490) client 0 local steps 0 loss 0.1766 train time 6.7182
(Trainer pid=1490) client 0 local steps 1 loss 0.1721 train time 6.6900
clientId: 0 current_loss: 0.16773484647274017 train_finish_times: [6.718162298202515, 6.68996524810791, 6.663710355758667]
(Trainer pid=1490) client 0 local steps 2 loss 0.1677 train time 6.6637
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8231310844421387 hit rate: 0.9145966172218323 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8231310844421387 hit rate: 0.9145966172218323 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8231310844421387 hit rate: 0.9145966172218323
global rounds: 6
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) Test AUC: 0.8231
(Trainer pid=1490) Test Hit Rate at 2: 0.9146
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1490) client 0 local steps 0 loss 0.1634 train time 6.6510
(Trainer pid=1490) client 0 local steps 1 loss 0.1592 train time 6.6749
clientId: 0 current_loss: 0.15507246553897858 train_finish_times: [6.650990724563599, 6.674895286560059, 6.665460586547852]
(Trainer pid=1490) client 0 local steps 2 loss 0.1551 train time 6.6655
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8247485756874084 hit rate: 0.9153063893318176 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8247485756874084 hit rate: 0.9153063893318176 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8247485756874084 hit rate: 0.9153063893318176
(Trainer pid=1490) Test AUC: 0.8247
(Trainer pid=1490) Test Hit Rate at 2: 0.9153
(Trainer pid=1490) Test Traveled User Hit Rate at 2: 1.0000
global rounds: 7
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1490) client 0 local steps 0 loss 0.1511 train time 6.6975
(Trainer pid=1490) client 0 local steps 1 loss 0.1474 train time 6.6825
clientId: 0 current_loss: 0.14376592636108398 train_finish_times: [6.697473764419556, 6.682470798492432, 6.71771764755249]
(Trainer pid=1490) client 0 local steps 2 loss 0.1438 train time 6.7177
//Log Theoretical Embedding Communication Cost Added (Train Phase): 594.87 MB //end
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) auc score: 0.8261364102363586 hit rate: 0.9166074991226196 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, b007f2c144f42677d65ffb9801000000) final auc score: 0.8261364102363586 hit rate: 0.9166074991226196 traveled user hit rate: 1.0
Predict Day 20 average auc score: 0.8261364102363586 hit rate: 0.9166074991226196
training is not complete
//train_time: 212301.278 ms//end
//Log Theoretical Pretrain Comm Cost: 0.00 MB //end
//Log Theoretical Train Comm Cost: 9519.91 MB //end
The whole process has ended
Total running time of the script: (4 minutes 1.489 seconds)