[Model] Implement Grace-POT (#171)

* Grace-POT

* Grace-POT

* [Model] Update Model

* nightly update

* [Model] Update Grace_POT

* nightly udpate

* nightly update

* nightly update

---------

Co-authored-by: dddg617 <996179900@qq.com>
Co-authored-by: dddg617 <75086617+dddg617@users.noreply.github.com>

Author: guofeng97

CONTRIBUTING.md

@@ -30,6 +30,7 @@ Contribution is always welcomed. Please feel free to open an issue or email to y
   - Xingyuan Ji (Gamma Lab)
   - Yuxin Guo (Gamma Lab)
   - Zihao Zhao (Gamma Lab)
+  - Feng Guo (Gamma Lab)
   - Yuxuan Shan (BUPT)
   - Zeyao Ma (BUPT)
   - Yiming Jia (BUPT)

README.md

@@ -406,6 +406,7 @@ CUDA_VISIBLE_DEVICES="1" TL_BACKEND="paddle" python gcn_trainer.py
 | [MVGRL [ICML 2020]](./examples/mvgrl)          | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: |           |
 | [InfoGraph [ICLR 2020]](./examples/infograph)  | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: |           |
 | [MERIT [IJCAI 2021]](./examples/merit)         | :heavy_check_mark: |                    | :heavy_check_mark: |           |
+| [GNN-POT [NeurIPS 2023]](./examples/grace_pot) |                    | :heavy_check_mark: |                    |           |
 
 | Heterogeneous Graph Learning                 | TensorFlow         | PyTorch            | Paddle             | MindSpore |
 | -------------------------------------------- | ------------------ | ------------------ | ------------------ | --------- |

examples/grace_pot/GRACE_POT_trainer.py

@@ -0,0 +1,290 @@
+import argparse
+import os.path as osp
+import os
+# os.environ['TL_BACKEND'] = 'torch'
+from time import perf_counter as t
+import yaml
+from yaml import SafeLoader
+import numpy as np
+import pickle
+import tensorlayerx as tlx
+from tensorlayerx.model import TrainOneStep, WithLoss
+from tensorlayerx.dataflow import random_split
+from gammagl.layers.conv import GCNConv
+from gammagl.datasets import Planetoid, Coauthor, Amazon
+import gammagl.transforms as T
+
+from gammagl.models.grace_pot import Grace_POT_Encoder, Grace_POT_Model
+from eval_gracepot import log_regression, MulticlassEvaluator
+
+A_upper_1 = None
+A_upper_2 = None
+A_lower_1 = None
+A_lower_2 = None
+
+class train_loss(WithLoss):
+    def __init__(self, model, drop_edge_rate_1, drop_edge_rate_2, use_pot=False, pot_batch=-1, kappa=0.5):
+        super(train_loss, self).__init__(backbone=model, loss_fn=None)
+        self.drop_edge_rate_1 = drop_edge_rate_1
+        self.drop_edge_rate_2 = drop_edge_rate_2
+        self.use_pot = use_pot
+        self.pot_batch = pot_batch
+        self.kappa = kappa
+
+    def forward(self, model, x, edge_index, epoch, data=None):
+        edge_index_1 = dropout_adj(edge_index, p=self.drop_edge_rate_1)[0]
+        edge_index_2 = dropout_adj(edge_index, p=self.drop_edge_rate_2)[0]
+        x_1, x_2 = x, x
+        z1 = model(x_1, edge_index_1)
+        z2 = model(x_2, edge_index_2)
+        node_list = np.arange(z1.shape[0])
+        np.random.shuffle(node_list)
+        
+        batch_size = 4096 if args.dataset in ["PubMed", "Computers", "WikiCS"] else None
+
+        if batch_size is not None:
+            node_list_batch = get_batch(node_list, batch_size, epoch)
+
+        # nce loss
+        if batch_size is not None:
+            z11 = z1[node_list_batch]
+            z22 = z2[node_list_batch]
+            nce_loss = model.loss(z11, z22)
+        else:
+            nce_loss = model.loss(z1, z2)
+
+        # pot loss
+        if self.use_pot:
+            # get node_list_tmp, the nodes to calculate pot_loss
+            if self.pot_batch != -1:
+                if batch_size is None:
+                    node_list_tmp = get_batch(node_list, self.pot_batch, epoch)
+                else:
+                    node_list_tmp = get_batch(node_list_batch, self.pot_batch, epoch)
+            else:
+                # full pot batch
+                if batch_size is None:
+                    node_list_tmp = node_list
+                else:
+                    node_list_tmp = node_list_batch
+                    
+            z11 = tlx.gather(z1, tlx.convert_to_tensor(node_list_tmp))
+            z22 = tlx.gather(z2, tlx.convert_to_tensor(node_list_tmp))
+
+            global A_upper_1, A_upper_2, A_lower_1, A_lower_2
+            if A_upper_1 is None or A_upper_2 is None:
+                A_upper_1, A_lower_1 = get_A_bounds(args.dataset, self.drop_edge_rate_1, args.cache)
+                A_upper_2, A_lower_2 = get_A_bounds(args.dataset, self.drop_edge_rate_2, args.cache)
+
+            pot_loss_1 = model.pot_loss(z11, z22, data.x, data.edge_index, edge_index_1, local_changes=self.drop_edge_rate_1, 
+                                          node_list=node_list_tmp, A_upper=A_upper_1, A_lower=A_lower_1)
+            pot_loss_2 = model.pot_loss(z22, z11, data.x, data.edge_index, edge_index_2, local_changes=self.drop_edge_rate_2, 
+                                          node_list=node_list_tmp, A_upper=A_upper_2, A_lower=A_lower_2)
+            pot_loss = (pot_loss_1 + pot_loss_2) / 2
+            loss = (1 - self.kappa) * nce_loss + self.kappa * pot_loss
+        else:
+            loss = nce_loss
+
+        return loss
+
+
+def test(model, data, dataset, split):
+    model.set_eval()
+    z = model(data.x, data.edge_index)
+    evaluator = MulticlassEvaluator()
+    res = log_regression(z, dataset, evaluator, split='preloaded', num_epochs=3000, preload_split=split)
+    return res
+
+def get_dataset(path, name):
+    assert name in ['Cora', 'CiteSeer', 'PubMed', 'Coauthor-CS', 'Coauthor-Phy', 'Computers', 'Photo']
+    name = 'dblp' if name == 'DBLP' else name
+
+    if name == 'Coauthor-CS':
+        return Coauthor(root=path, name='cs', transform=T.NormalizeFeatures())
+
+    if name == 'Coauthor-Phy':
+        return Coauthor(root=path, name='physics', transform=T.NormalizeFeatures())
+
+    if name == 'Computers':
+        return Amazon(root=path, name='computers', transform=T.NormalizeFeatures())
+
+    if name == 'Photo':
+        return Amazon(root=path, name='photo', transform=T.NormalizeFeatures())
+
+
+    return (Planetoid)(path, name, transform=T.NormalizeFeatures()) # public split
+
+def generate_split(num_samples: int, train_ratio: float, val_ratio: float):
+    train_len = int(num_samples * train_ratio)
+    val_len = int(num_samples * val_ratio)
+    test_len = num_samples - train_len - val_len
+
+    train_set, test_set, val_set = random_split(tlx.arange(0, num_samples), (train_len, test_len, val_len))
+
+    idx_train, idx_test, idx_val = train_set.indices, test_set.indices, val_set.indices
+    train_mask = tlx.zeros((num_samples,)).to(tlx.bool)
+    test_mask = tlx.zeros((num_samples,)).to(tlx.bool)
+    val_mask = tlx.zeros((num_samples,)).to(tlx.bool)
+
+    train_mask[idx_train] = True
+    test_mask[idx_test] = True
+    val_mask[idx_val] = True
+
+    return train_mask, test_mask, val_mask
+
+def get_batch(node_list, batch_size, epoch):
+    num_nodes = len(node_list)
+    num_batches = (num_nodes - 1) // batch_size + 1
+    i = epoch % num_batches
+    if (i + 1) * batch_size >= len(node_list):
+        node_list_batch = node_list[i * batch_size:]
+    else:
+        node_list_batch = node_list[i * batch_size:(i + 1) * batch_size]
+    return node_list_batch
+
+def get_A_bounds(dataset, drop_rate, cache):
+    upper_lower_file = osp.join(cache, f"{dataset}_{drop_rate}_upper_lower.pkl")
+    if osp.exists(upper_lower_file):
+        with open(upper_lower_file, 'rb') as file:       
+            A_upper, A_lower = pickle.load(file)
+    else:
+        A_upper, A_lower = None, None
+    return A_upper, A_lower
+
+def filter_adj(row, col, edge_attr, mask):
+    mask = tlx.convert_to_tensor(mask, dtype=tlx.bool)
+    return row[mask], col[mask], None if edge_attr is None else edge_attr[mask]
+
+def dropout_adj(
+    edge_index,
+    edge_attr = None,
+    p = 0.5,
+    force_undirected = False,
+    num_nodes = None,
+    training = True,
+):
+
+    if p < 0. or p > 1.:
+        raise ValueError(f'Dropout probability has to be between 0 and 1 '
+                         f'(got {p}')
+
+    if not training or p == 0.0:
+        return edge_index, edge_attr
+
+    # row, col = edge_index
+    row = edge_index[0]
+    col = edge_index[1]
+
+    mask = np.random.random(tlx.get_tensor_shape(row)) >= p
+
+    if force_undirected:
+        mask[row > col] = False
+
+    row, col, edge_attr = filter_adj(row, col, edge_attr, mask)
+
+    if force_undirected:
+        edge_index = tlx.stack(
+            [tlx.concat([row, col], 0),
+             tlx.concat([col, row], 0)], dim=0)
+        if edge_attr is not None:
+            edge_attr = tlx.concat([edge_attr, edge_attr], 0)
+    else:
+        edge_index = tlx.stack([row, col])
+
+    return edge_index, edge_attr
+
+
+def main(args):
+    if args.gpu_id >= 0:
+        tlx.set_device(device='GPU', id=args.gpu_id)
+    else:
+        tlx.set_device(device='CPU')
+
+    config = yaml.load(open(args.config), Loader=SafeLoader)[args.dataset]
+    # for hyperparameter tuning
+    if args.drop_1 != -1:
+        config['drop_edge_rate_1'] = args.drop_1
+    if args.drop_2 != -1:
+        config['drop_edge_rate_2'] = args.drop_2
+    if args.tau != -1:
+        config['tau'] = args.tau
+    if args.num_epochs != -1:
+        config['num_epochs'] = args.num_epochs
+    print(args)
+    print(config)
+
+    learning_rate = config['learning_rate']
+    num_hidden = config['num_hidden']
+    num_proj_hidden = config['num_proj_hidden']
+    activation = ({'relu': tlx.nn.ReLU, 'prelu': tlx.nn.PRelu()})[config['activation']]
+    base_model = ({'GCNConv': GCNConv})[config['base_model']]
+    num_layers = config['num_layers']
+
+    drop_edge_rate_1 = config['drop_edge_rate_1']
+    drop_edge_rate_2 = config['drop_edge_rate_2']
+    tau = config['tau']
+    num_epochs = config['num_epochs']
+    weight_decay = config['weight_decay']
+    use_pot = args.use_pot
+    kappa = args.kappa
+    pot_batch = args.pot_batch
+
+    dataset = get_dataset(args.path, args.dataset)
+    data = dataset[0]
+    
+    # generate split
+    if args.dataset in ["Cora", "CiteSeer", "PubMed"]:
+        split = data.train_mask, data.val_mask, data.test_mask
+        print("Public Split")
+    else:
+        split = generate_split(data.num_nodes, train_ratio=0.1, val_ratio=0.1)
+        print("Random Split")
+
+    encoder = Grace_POT_Encoder(dataset.num_features, num_hidden, activation,
+                      base_model=base_model, k=num_layers)
+    model = Grace_POT_Model(encoder, num_hidden, num_proj_hidden, tau, dataset=args.dataset, cached=args.cache)
+    train_weights = model.trainable_weights
+    optimizer = tlx.optimizers.Adam(lr=learning_rate, weight_decay=weight_decay)
+    loss_func = train_loss(model, drop_edge_rate_1, drop_edge_rate_2, use_pot, pot_batch, kappa)
+    train_one_step = TrainOneStep(loss_func, optimizer, train_weights)
+
+    #timing        
+    start = t()
+    prev = start
+    for epoch in range(1, num_epochs + 1):
+        model.set_train()
+        loss=train_one_step(model, data.x, data.edge_index, epoch ,data)
+        now = t()
+        print(f'(T) | Epoch={epoch:03d}, loss={loss:.4f}, '
+              f'this epoch {now - prev:.4f}, total {now - start:.4f}')
+        if epoch % 100 == 0:
+            res = test(model, data, dataset, split)
+            print(res)
+        prev = now
+
+    print("=== Final ===")
+    res = test(model, data, dataset, split)
+    print(res)
+
+
+if __name__ == '__main__':
+    # parameters setting
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--path', type=str, default="./")
+    parser.add_argument('--cache', type=str, default="./")
+    parser.add_argument('--dataset', type=str, default='Cora')
+    parser.add_argument('--gpu_id', type=int, default=0)
+    parser.add_argument('--config', type=str, default='./config.yaml')
+    parser.add_argument('--use_pot', default=True) # whether to use pot in loss
+    parser.add_argument('--kappa', type=float, default=0.5)
+    parser.add_argument('--pot_batch', type=int, default=-1)
+    parser.add_argument('--drop_1', type=float, default=0.4)
+    parser.add_argument('--drop_2', type=float, default=0.3)
+    parser.add_argument('--tau', type=float, default=0.9) # temperature of nce loss
+    parser.add_argument('--num_epochs',type=int,default=-1)
+    parser.add_argument('--save_file', type=str, default=".")
+    parser.add_argument('--seed', type=int, default=12345)
+    args = parser.parse_args()
+    main(args)
+
+    

examples/grace_pot/README.md

@@ -0,0 +1,49 @@
+# GammaGL Implementation of GRACE-POT
+This GammaGL example implements the model proposed in the paper [xxxx](https://arxiv.org/).
+
+Author's code:
+
+## Example Implementor
+
+This example was implemented by Siyuan Zhang
+
+## Datasets
+
+##### Unsupervised Node Classification Datasets:
+
+'Cora', 'Citeseer' and 'Pubmed'
+
+| Dataset  | # Nodes | # Edges | # Classes |
+| -------- | ------- | ------- | --------- |
+| Cora     | 2,708   | 10,556  | 7         |
+| Pubmed   | 19,717  | 88,651  | 3         |
+| Photo    | 7,650   | 238,162 | 8         |
+
+
+## How to run examples
+Fisrt, make the directories for datasets and bounds to save
+``` bash
+mkdir ~/datasets
+mkdir ~/datasets/bounds
+```
+Then, go into the directory of a model. If you want to set the parameters, you should modify the ocnfiguration files in the directory ("config.yaml" for GRACE). The following is the command line to run each model (dataset used is Cora for example):
+```bash
+# original GRACE
+python GRACE_POT_trainer.py --dataset Cora --gpu_id 0
+# GRACE + POT
+python GRACE_POT_trainer.py --dataset Cora --gpu_id 0 --use_pot --kappa 0.4
+```
+The result will be appended to the file "res/{dataset_name}_base_temp.csv" and "res/{dataset_name}_pot_temp.csv" respectively. You can also set the parameter "save_file" to specify the file to save results. We use minibatch to reduce the memory occupation, you can modify it in the code. 
+
+## 	Performance
+
+```
+                    |   Author's Code   |  GAMMAGL's Code   |
+| Dataset | Metrics |  GRACE  |GRACE-POT|  GRACE  |GRACE-POT|
+|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|
+|  Cora   |  Mi-F1  |   78.2  |   79.2  |   78.2  |   82.2  |
+|  Cora   |  Ma-F1  |   76.8  |   77.8  |   77.1  |   81.3  |
+|  PubMed |  Mi-F1  |   81.6  |   82.0  |   81.6  |   82.0  |
+|  PubMed |  Ma-F1  |   81.7  |   82.4  |   80.5  |   80.1  |
+|  Photo  |  Mi-F1  |   91.2  |   91.8  |   89.8  |   90.0  |
+|  Photo  |  Ma-F1  |   89.2  |   90.0  |   88.5  |   87.9  |