wehub-resource-sync
66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
import argparse
|
|
|
|
import dgl
|
|
|
|
import torch
|
|
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
|
|
from dgl.nn import LabelPropagation
|
|
|
|
|
|
def main():
|
|
# check cuda
|
|
device = (
|
|
f"cuda:{args.gpu}"
|
|
if torch.cuda.is_available() and args.gpu >= 0
|
|
else "cpu"
|
|
)
|
|
|
|
# load data
|
|
if args.dataset == "Cora":
|
|
dataset = CoraGraphDataset()
|
|
elif args.dataset == "Citeseer":
|
|
dataset = CiteseerGraphDataset()
|
|
elif args.dataset == "Pubmed":
|
|
dataset = PubmedGraphDataset()
|
|
else:
|
|
raise ValueError("Dataset {} is invalid.".format(args.dataset))
|
|
|
|
g = dataset[0]
|
|
g = dgl.add_self_loop(g)
|
|
|
|
labels = g.ndata.pop("label").to(device).long()
|
|
|
|
# load masks for train / test, valid is not used.
|
|
train_mask = g.ndata.pop("train_mask")
|
|
test_mask = g.ndata.pop("test_mask")
|
|
|
|
train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze().to(device)
|
|
test_idx = torch.nonzero(test_mask, as_tuple=False).squeeze().to(device)
|
|
|
|
g = g.to(device)
|
|
|
|
# label propagation
|
|
lp = LabelPropagation(args.num_layers, args.alpha)
|
|
logits = lp(g, labels, mask=train_idx)
|
|
|
|
test_acc = torch.sum(
|
|
logits[test_idx].argmax(dim=1) == labels[test_idx]
|
|
).item() / len(test_idx)
|
|
print("Test Acc {:.4f}".format(test_acc))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""
|
|
Label Propagation Hyperparameters
|
|
"""
|
|
parser = argparse.ArgumentParser(description="LP")
|
|
parser.add_argument("--gpu", type=int, default=0)
|
|
parser.add_argument("--dataset", type=str, default="Cora")
|
|
parser.add_argument("--num-layers", type=int, default=10)
|
|
parser.add_argument("--alpha", type=float, default=0.5)
|
|
|
|
args = parser.parse_args()
|
|
print(args)
|
|
|
|
main()
|