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]
100%|██████████| 552k/552k [00:00<00:00, 12.1MB/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=ae0bffd2-f7f6-4977-b5f1-fd2ef356a833
To: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_global.txt

  0%|          | 0.00/278M [00:00<?, ?B/s]
  1%|          | 1.57M/278M [00:00<00:17, 15.7MB/s]
  3%|▎         | 8.91M/278M [00:00<00:12, 22.0MB/s]
  6%|▌         | 17.3M/278M [00:00<00:07, 36.9MB/s]
 10%|▉         | 27.8M/278M [00:00<00:04, 54.9MB/s]
 12%|█▏        | 34.6M/278M [00:00<00:05, 46.7MB/s]
 15%|█▌        | 42.5M/278M [00:00<00:04, 48.5MB/s]
 19%|█▉        | 53.0M/278M [00:01<00:03, 61.5MB/s]
 23%|██▎       | 64.0M/278M [00:01<00:02, 73.2MB/s]
 26%|██▌       | 72.4M/278M [00:01<00:03, 59.6MB/s]
 29%|██▊       | 79.7M/278M [00:01<00:03, 54.8MB/s]
 35%|███▍      | 95.9M/278M [00:01<00:02, 78.5MB/s]
 39%|███▉      | 110M/278M [00:01<00:01, 88.3MB/s]
 45%|████▌     | 126M/278M [00:01<00:01, 90.6MB/s]
 51%|█████▏    | 143M/278M [00:02<00:01, 94.2MB/s]
 55%|█████▌    | 153M/278M [00:02<00:01, 69.6MB/s]
 58%|█████▊    | 161M/278M [00:02<00:01, 70.9MB/s]
 64%|██████▎   | 177M/278M [00:02<00:01, 80.5MB/s]
 69%|██████▉   | 191M/278M [00:02<00:00, 94.3MB/s]
 75%|███████▌  | 210M/278M [00:02<00:00, 115MB/s]
 82%|████████▏ | 228M/278M [00:02<00:00, 132MB/s]
 89%|████████▉ | 247M/278M [00:03<00:00, 147MB/s]
 96%|█████████▌| 267M/278M [00:03<00:00, 160MB/s]
100%|██████████| 278M/278M [00:03<00:00, 87.0MB/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]
 26%|██▌       | 7.34M/28.7M [00:00<00:00, 72.9MB/s]
 86%|████████▌ | 24.6M/28.7M [00:00<00:00, 130MB/s]
100%|██████████| 28.7M/28.7M [00:00<00:00, 127MB/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=1011) checking code and file path: JP,/home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset
(Trainer pid=1011) printing in getdata, path: /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset
(Trainer pid=1011) Loading data in /home/docs/checkouts/readthedocs.org/user_builds/fedgraph/checkouts/latest/data/LPDataset/data_JP.txt
(Trainer pid=1011) Device: 'cpu'
(Trainer pid=1011) loading train_data and test_data
(Trainer pid=1011) client 0 local steps 0 loss 0.7621 train time 6.1753
(Trainer pid=1011) client 0 local steps 1 loss 0.6571 train time 6.0053
clientId: 0 current_loss: 0.5708792209625244 train_finish_times: [6.17525577545166, 6.0052571296691895, 5.833333253860474]
(Trainer pid=1011) client 0 local steps 2 loss 0.5709 train time 5.8333
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.7258461117744446 hit rate: 0.8547433018684387 traveled user hit rate: 0.8571428656578064
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.7258461117744446 hit rate: 0.8547433018684387 traveled user hit rate: 0.8571428656578064

Predict Day 20 average auc score: 0.7258461117744446 hit rate: 0.8547433018684387
global rounds: 1
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) Test AUC: 0.7258
(Trainer pid=1011) Test Hit Rate at 2: 0.8547
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 0.8571
(Trainer pid=1011) client 0 local steps 0 loss 0.4959 train time 5.9317
(Trainer pid=1011) client 0 local steps 1 loss 0.4305 train time 5.8241
clientId: 0 current_loss: 0.3744064271450043 train_finish_times: [5.931670188903809, 5.824119329452515, 5.804453372955322]
(Trainer pid=1011) client 0 local steps 2 loss 0.3744 train time 5.8045
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.8070921301841736 hit rate: 0.9063165187835693 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.8070921301841736 hit rate: 0.9063165187835693 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8070921301841736 hit rate: 0.9063165187835693
global rounds: 2
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) Test AUC: 0.8071
(Trainer pid=1011) Test Hit Rate at 2: 0.9063
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1011) client 0 local steps 0 loss 0.3270 train time 5.8339
(Trainer pid=1011) client 0 local steps 1 loss 0.2876 train time 5.6701
clientId: 0 current_loss: 0.2559906840324402 train_finish_times: [5.833918571472168, 5.670072555541992, 5.840089797973633]
(Trainer pid=1011) client 0 local steps 2 loss 0.2560 train time 5.8401
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.817216157913208 hit rate: 0.9123491644859314 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.817216157913208 hit rate: 0.9123491644859314 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.817216157913208 hit rate: 0.9123491644859314
global rounds: 3
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) Test AUC: 0.8172
(Trainer pid=1011) Test Hit Rate at 2: 0.9123
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1011) client 0 local steps 0 loss 0.2318 train time 5.8853
(Trainer pid=1011) client 0 local steps 1 loss 0.2141 train time 5.7716
clientId: 0 current_loss: 0.20163379609584808 train_finish_times: [5.885336399078369, 5.771606206893921, 5.795802354812622]
(Trainer pid=1011) client 0 local steps 2 loss 0.2016 train time 5.7958
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.8199998736381531 hit rate: 0.9112846255302429 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.8199998736381531 hit rate: 0.9112846255302429 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8199998736381531 hit rate: 0.9112846255302429
global rounds: 4
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) Test AUC: 0.8200
(Trainer pid=1011) Test Hit Rate at 2: 0.9113
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1011) client 0 local steps 0 loss 0.1927 train time 5.8475
(Trainer pid=1011) client 0 local steps 1 loss 0.1861 train time 5.7994
clientId: 0 current_loss: 0.1807977706193924 train_finish_times: [5.847520589828491, 5.799439430236816, 5.9108757972717285]
(Trainer pid=1011) client 0 local steps 2 loss 0.1808 train time 5.9109
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.8219790458679199 hit rate: 0.9128223061561584 traveled user hit rate: 1.0
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.8219790458679199 hit rate: 0.9128223061561584 traveled user hit rate: 1.0

Predict Day 20 average auc score: 0.8219790458679199 hit rate: 0.9128223061561584
global rounds: 5
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) Test AUC: 0.8220
(Trainer pid=1011) Test Hit Rate at 2: 0.9128
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 1.0000
(Trainer pid=1011) client 0 local steps 0 loss 0.1761 train time 5.8432
(Trainer pid=1011) client 0 local steps 1 loss 0.1716 train time 5.8751
clientId: 0 current_loss: 0.16727741062641144 train_finish_times: [5.843219518661499, 5.875112056732178, 5.796710729598999]
(Trainer pid=1011) client 0 local steps 2 loss 0.1673 train time 5.7967
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.8238192796707153 hit rate: 0.9137686491012573 traveled user hit rate: 0.8571428656578064
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.8238192796707153 hit rate: 0.9137686491012573 traveled user hit rate: 0.8571428656578064

Predict Day 20 average auc score: 0.8238192796707153 hit rate: 0.9137686491012573
(Trainer pid=1011) Test AUC: 0.8238
(Trainer pid=1011) Test Hit Rate at 2: 0.9138
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 0.8571
global rounds: 6
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) client 0 local steps 0 loss 0.1630 train time 5.9148
(Trainer pid=1011) client 0 local steps 1 loss 0.1588 train time 5.8154
clientId: 0 current_loss: 0.15470828115940094 train_finish_times: [5.914848327636719, 5.815420389175415, 5.881413459777832]
(Trainer pid=1011) client 0 local steps 2 loss 0.1547 train time 5.8814
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.8254882097244263 hit rate: 0.9148331880569458 traveled user hit rate: 0.8571428656578064
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.8254882097244263 hit rate: 0.9148331880569458 traveled user hit rate: 0.8571428656578064

Predict Day 20 average auc score: 0.8254882097244263 hit rate: 0.9148331880569458
global rounds: 7
Training in LP_train_global_round, number of clients: 1
(Trainer pid=1011) Test AUC: 0.8255
(Trainer pid=1011) Test Hit Rate at 2: 0.9148
(Trainer pid=1011) Test Traveled User Hit Rate at 2: 0.8571
(Trainer pid=1011) client 0 local steps 0 loss 0.1508 train time 5.8490
(Trainer pid=1011) client 0 local steps 1 loss 0.1470 train time 5.9018
clientId: 0 current_loss: 0.14346961677074432 train_finish_times: [5.849045038223267, 5.901822566986084, 5.846710443496704]
(Trainer pid=1011) client 0 local steps 2 loss 0.1435 train time 5.8467
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) auc score: 0.8269479870796204 hit rate: 0.9171989560127258 traveled user hit rate: 0.8571428656578064
Day 0 client Actor(run_LP.<locals>.setup_trainer_server.<locals>.Trainer, f21cad615771f90ffb8dd16601000000) final auc score: 0.8269479870796204 hit rate: 0.9171989560127258 traveled user hit rate: 0.8571428656578064

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

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

Gallery generated by Sphinx-Gallery