import torch
import torch.nn.functional as F
from torch.nn import ModuleList
from tqdm import tqdm
from torch_geometric.datasets import Reddit
from torch_geometric.loader import ClusterData, ClusterLoader, NeighborSampler
from torch_geometric.nn import SAGEConv
dataset = Reddit('../data/Reddit')
data = dataset[0]
cluster_data = ClusterData(data,
num_parts=1500,
recursive=False,
save_dir=dataset.processed_dir)
train_loader = ClusterLoader(cluster_data,
batch_size=20,
shuffle=True,
num_workers=12)
subgraph_loader = NeighborSampler(data.edge_index,
sizes=[-1],
batch_size=1024,
shuffle=False,
num_workers=12)
class Net(torch.nn.Module):
def __init__(self, in_channels, out_channels):
super(Net, self).__init__()
self.convs = ModuleList(
import torch.nn.functional as F
from torch_geometric.datasets import PPI
from torch_geometric.nn import SAGEConv, BatchNorm
from torch_geometric.data import Batch
from torch_geometric.loader import DataLoader
from torch_geometric.loader import ClusterData, ClusterLoader
from sklearn.metrics import f1_score
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'PPI')
train_dataset = PPI(path, split='train')
val_dataset = PPI(path, split='val')
test_dataset = PPI(path, split='test')
train_data = Batch.from_data_list(train_dataset)
cluster_data = ClusterData(train_data,
num_parts=50,
recursive=False,
save_dir=train_dataset.processed_dir)
train_loader = ClusterLoader(cluster_data,
batch_size=1,
shuffle=True,
num_workers=12)
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)
class Net(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
super(Net, self).__init__()
self.convs = torch.nn.ModuleList()
self.batch_norms = torch.nn.ModuleList()
def test_cluster_gcn():
adj = torch.tensor([
[1, 1, 1, 0, 1, 0],
[1, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
])
x = torch.Tensor([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
edge_index = adj.nonzero(as_tuple=False).t()
edge_attr = torch.arange(edge_index.size(1))
data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
data.num_nodes = 6
cluster_data = ClusterData(data, num_parts=2, log=False)
assert cluster_data.partptr.tolist() == [0, 3, 6]
assert cluster_data.perm.tolist() == [0, 2, 4, 1, 3, 5]
assert cluster_data.data.x.tolist() == [
[0, 0],
[2, 2],
[4, 4],
[1, 1],
[3, 3],
[5, 5],
]
assert cluster_data.data.adj.to_dense().tolist() == [
[0, 2, 3, 1, 0, 0],
[8, 9, 10, 0, 0, 0],
[14, 15, 16, 0, 0, 0],
[4, 0, 0, 5, 6, 7],
[0, 0, 0, 11, 12, 13],
[0, 0, 0, 17, 18, 19],
]
data = cluster_data[0]
assert data.num_nodes == 3
assert data.x.tolist() == [[0, 0], [2, 2], [4, 4]]
assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
[0, 1, 2, 0, 1, 2, 0, 1, 2]]
assert data.edge_attr.tolist() == [0, 2, 3, 8, 9, 10, 14, 15, 16]
data = cluster_data[1]
assert data.num_nodes == 3
assert data.x.tolist() == [[1, 1], [3, 3], [5, 5]]
assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
[0, 1, 2, 0, 1, 2, 0, 1, 2]]
assert data.edge_attr.tolist() == [5, 6, 7, 11, 12, 13, 17, 18, 19]
loader = ClusterLoader(cluster_data, batch_size=1)
iterator = iter(loader)
data = next(iterator)
assert data.x.tolist() == [[0, 0], [2, 2], [4, 4]]
assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
[0, 1, 2, 0, 1, 2, 0, 1, 2]]
assert data.edge_attr.tolist() == [0, 2, 3, 8, 9, 10, 14, 15, 16]
data = next(iterator)
assert data.x.tolist() == [[1, 1], [3, 3], [5, 5]]
assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
[0, 1, 2, 0, 1, 2, 0, 1, 2]]
assert data.edge_attr.tolist() == [5, 6, 7, 11, 12, 13, 17, 18, 19]
loader = ClusterLoader(cluster_data, batch_size=2, shuffle=False)
data = next(iter(loader))
assert data.num_nodes == 6
assert data.x.tolist() == [
[0, 0],
[2, 2],
[4, 4],
[1, 1],
[3, 3],
[5, 5],
]
assert to_dense_adj(data.edge_index).squeeze().tolist() == [
[1, 1, 1, 1, 0, 0],
[1, 1, 1, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[1, 0, 0, 1, 1, 1],
[0, 0, 0, 1, 1, 1],
[0, 0, 0, 1, 1, 1],
]