1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
"""Dataset and graph-operator utilities for the MVP."""
from __future__ import annotations
from pathlib import Path
from typing import Any
from urllib.request import urlopen
import torch
import torch_geometric.transforms as transforms
from torch_geometric.datasets import Planetoid
PLANETOID_FILES = (
"x",
"tx",
"allx",
"y",
"ty",
"ally",
"graph",
"test.index",
)
def _ensure_planetoid_raw(root: str | Path, name: str) -> None:
"""Pre-fetch from raw.githubusercontent.com to avoid a flaky redirect."""
raw_dir = Path(root) / name / "raw"
raw_dir.mkdir(parents=True, exist_ok=True)
lower_name = name.lower()
base_url = (
"https://raw.githubusercontent.com/kimiyoung/"
"planetoid/master/data"
)
for suffix in PLANETOID_FILES:
filename = f"ind.{lower_name}.{suffix}"
destination = raw_dir / filename
if destination.exists() and destination.stat().st_size > 0:
continue
temporary = destination.with_suffix(destination.suffix + ".part")
with urlopen(f"{base_url}/{filename}", timeout=60) as response:
temporary.write_bytes(response.read())
temporary.replace(destination)
def spmm(matrix: torch.Tensor, dense: torch.Tensor) -> torch.Tensor:
"""Multiply a native PyTorch sparse matrix by a dense matrix."""
return torch.sparse.mm(matrix, dense)
def normalized_adjacency(
edge_index: torch.Tensor,
num_nodes: int,
) -> torch.Tensor:
"""Return A-hat = D^-1/2 (A + I) D^-1/2 as a coalesced COO tensor."""
row, col = edge_index
loops = torch.arange(num_nodes, device=edge_index.device)
row = torch.cat((row, loops))
col = torch.cat((col, loops))
degree = torch.zeros(num_nodes, device=edge_index.device)
degree.scatter_add_(0, row, torch.ones_like(row, dtype=torch.float32))
inverse_sqrt = degree.pow(-0.5)
inverse_sqrt.masked_fill_(torch.isinf(inverse_sqrt), 0.0)
values = inverse_sqrt[row] * inverse_sqrt[col]
return torch.sparse_coo_tensor(
torch.stack((row, col)),
values,
(num_nodes, num_nodes),
).coalesce()
def load_cora(
root: str | Path = "data",
device: str | torch.device = "cpu",
) -> dict[str, Any]:
"""Download Cora and return the tensors used by both training rules."""
_ensure_planetoid_raw(root, "Cora")
dataset = Planetoid(
root=str(root),
name="Cora",
transform=transforms.NormalizeFeatures(),
)
graph = dataset[0]
adjacency = normalized_adjacency(
graph.edge_index,
graph.num_nodes,
).to(device)
return {
"x": graph.x.to(device),
"y": graph.y.to(device),
"adjacency": adjacency,
"train_mask": graph.train_mask.to(device),
"val_mask": graph.val_mask.to(device),
"test_mask": graph.test_mask.to(device),
"num_nodes": graph.num_nodes,
"num_features": dataset.num_features,
"num_classes": dataset.num_classes,
}
|
