def test_pna_conv():
x = torch.randn(4, 16)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
row, col = edge_index
value = torch.rand(row.size(0), 3)
adj = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
conv = PNAConv(16,
32,
aggregators,
scalers,
deg=torch.tensor([0, 3, 0, 1]),
edge_dim=3,
towers=4)
assert conv.__repr__() == 'PNAConv(16, 32, towers=4, edge_dim=3)'
out = conv(x, edge_index, value)
assert out.size() == (4, 32)
assert torch.allclose(conv(x, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x, edge_index, value), out, atol=1e-6)
t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x, adj.t()), out, atol=1e-6)
def test_pna_conv_get_degree_histogram():
edge_index = torch.tensor([[0, 0, 0, 1, 1, 2, 3], [1, 2, 3, 2, 0, 0, 0]])
data = Data(num_nodes=5, edge_index=edge_index)
loader = NeighborLoader(
data,
num_neighbors=[-1],
input_nodes=None,
batch_size=5,
shuffle=False,
)
deg_hist = PNAConv.get_degree_histogram(loader)
deg_hist_ref = torch.tensor([1, 2, 1, 1])
assert torch.equal(deg_hist_ref, deg_hist)
edge_index_1 = torch.tensor([[0, 0, 0, 1, 1, 2, 3], [1, 2, 3, 2, 0, 0, 0]])
edge_index_2 = torch.tensor([[1, 1, 2, 2, 0, 3, 3], [2, 3, 3, 1, 1, 0, 2]])
edge_index_3 = torch.tensor([[1, 3, 2, 0, 0, 4, 2], [2, 0, 4, 1, 1, 0, 3]])
edge_index_4 = torch.tensor([[0, 1, 2, 4, 0, 1, 3], [2, 3, 3, 1, 1, 0, 2]])
data_1 = Data(num_nodes=5,
edge_index=edge_index_1) # deg_hist = [1, 2 ,1 ,1]
data_2 = Data(num_nodes=5, edge_index=edge_index_2) # deg_hist = [1, 1, 3]
data_3 = Data(num_nodes=5, edge_index=edge_index_3) # deg_hist = [0, 3, 2]
data_4 = Data(num_nodes=5, edge_index=edge_index_4) # deg_hist = [1, 1, 3]
loader = DataLoader(
[data_1, data_2, data_3, data_4],
batch_size=1,
shuffle=False,
)
deg_hist = PNAConv.get_degree_histogram(loader)
deg_hist_ref = torch.tensor([3, 7, 9, 1])
assert torch.equal(deg_hist_ref, deg_hist)
def __init__(self):
super(Net, self).__init__()
aggregators = ['mean', 'min', 'max', 'std']
scalers = ['identity', 'amplification', 'attenuation']
self.dropout = args.dropout
self.patience = args.patience
self.convs = ModuleList()
self.convs.append(
PNAConv(in_channels=dataset.num_features,
out_channels=args.hidden,
aggregators=aggregators,
scalers=scalers,
deg=deg,
edge_dim=dataset.num_edge_features,
towers=1,
pre_layers=args.pretrans_layers,
post_layers=args.posttrans_layers,
divide_input=False))
for _ in range(args.n_conv_layers):
conv = PNAConv(in_channels=args.hidden,
out_channels=args.hidden,
aggregators=aggregators,
scalers=scalers,
deg=deg,
edge_dim=dataset.num_edge_features,
towers=args.towers,
pre_layers=args.pretrans_layers,
post_layers=args.posttrans_layers,
divide_input=False)
self.convs.append(conv)
gr = []
g = list(
map(
int,
np.ceil(
np.geomspace(args.hidden, dataset.num_classes,
args.mlp_layers + 1))))
g[0] = args.hidden
g[-1] = dataset.num_classes
for i in range(args.mlp_layers):
gr.append(Linear(g[i], g[i + 1]))
if i < args.mlp_layers - 1:
gr.append(Dropout(p=self.dropout))
gr.append(LogSoftmax() if i == args.mlp_layers - 1 else ReLU())
self.mlp = Sequential(*gr)
def __init__(self):
super().__init__()
self.node_emb = Embedding(21, 75)
self.edge_emb = Embedding(4, 50)
aggregators = ['mean', 'min', 'max', 'std']
scalers = ['identity', 'amplification', 'attenuation']
self.convs = ModuleList()
self.batch_norms = ModuleList()
for _ in range(4):
conv = PNAConv(in_channels=75,
out_channels=75,
aggregators=aggregators,
scalers=scalers,
deg=deg,
edge_dim=50,
towers=5,
pre_layers=1,
post_layers=1,
divide_input=False)
self.convs.append(conv)
self.batch_norms.append(BatchNorm(75))
self.mlp = Sequential(Linear(75, 50), ReLU(), Linear(50, 25), ReLU(),
Linear(25, 1))
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.save_hyperparameters()
kwargs = self.sanetize_kwargs(kwargs)
self.node_emb = Embedding(kwargs["node_vocab"], kwargs["node_dim"])
self.edge_emb = Embedding(kwargs["edge_vocab"], kwargs["edge_dim"])
self.convs = ModuleList()
self.batch_norms = ModuleList()
for _ in range(kwargs["num_layers"]):
conv = PNAConv(
in_channels=kwargs["node_dim"],
out_channels=kwargs["node_dim"],
aggregators=kwargs["aggregators"],
scalers=kwargs["scalers"],
deg=torch.tensor(kwargs["deg"]),
edge_dim=kwargs["edge_dim"],
towers=kwargs["towers"],
pre_layers=kwargs["pre_layers"],
post_layers=kwargs["post_layers"],
divide_input=kwargs["divide_input"],
)
self.convs.append(conv)
self.batch_norms.append(BatchNorm(kwargs["node_dim"]))
self.mlp = Sequential(
Linear(kwargs["node_dim"], kwargs["edge_dim"]),
ReLU(),
Linear(kwargs["edge_dim"], kwargs["hidden_channels"]),
ReLU(),
Linear(kwargs["hidden_channels"], kwargs["num_classes"]),
)
def make_graph_layer(self, hidden_dim, layer_idx):
return PNAConv(
hidden_dim,
hidden_dim,
aggregators=["mean", "min", "max", "std"],
scalers=["identity", "amplification", "attenuation"],
deg=self.deg,
towers=4,
divide_input=True,
)
def run(args: argparse.ArgumentParser) -> None:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('BENCHMARK STARTS')
for dataset_name in args.datasets:
assert dataset_name in supported_sets.keys(
), f"Dataset {dataset_name} isn't supported."
print(f'Dataset: {dataset_name}')
dataset, num_classes = get_dataset(dataset_name, args.root,
args.use_sparse_tensor)
data = dataset.to(device)
hetero = True if dataset_name == 'ogbn-mag' else False
mask = ('paper', None) if dataset_name == 'ogbn-mag' else None
degree = None
inputs_channels = data[
'paper'].num_features if dataset_name == 'ogbn-mag' \
else dataset.num_features
for model_name in args.models:
if model_name not in supported_sets[dataset_name]:
print(f'Configuration of {dataset_name} + {model_name} '
f'not supported. Skipping.')
continue
print(f'Evaluation bench for {model_name}:')
for batch_size in args.eval_batch_sizes:
if not hetero:
subgraph_loader = NeighborLoader(
data,
num_neighbors=[-1], # layer-wise inference
input_nodes=mask,
batch_size=batch_size,
shuffle=False,
num_workers=args.num_workers,
)
for layers in args.num_layers:
if hetero:
subgraph_loader = NeighborLoader(
data,
num_neighbors=[args.hetero_num_neighbors] *
layers, # batch-wise inference
input_nodes=mask,
batch_size=batch_size,
shuffle=False,
num_workers=args.num_workers,
)
for hidden_channels in args.num_hidden_channels:
print('----------------------------------------------')
print(
f'Batch size={batch_size}, '
f'Layers amount={layers}, '
f'Num_neighbors={subgraph_loader.num_neighbors}, '
f'Hidden features size={hidden_channels}')
params = {
'inputs_channels': inputs_channels,
'hidden_channels': hidden_channels,
'output_channels': num_classes,
'num_heads': args.num_heads,
'num_layers': layers,
}
if model_name == 'pna':
if degree is None:
degree = PNAConv.get_degree_histogram(
subgraph_loader)
print(f'Calculated degree for {dataset_name}.')
params['degree'] = degree
model = get_model(
model_name, params,
metadata=data.metadata() if hetero else None)
model = model.to(device)
model.eval()
for _ in range(args.warmup):
model.inference(subgraph_loader, device,
progress_bar=True)
if args.experimental_mode:
with torch_geometric.experimental_mode():
with timeit():
model.inference(subgraph_loader, device,
progress_bar=True)
else:
with timeit():
model.inference(subgraph_loader, device,
progress_bar=True)
if args.profile:
with torch_profile():
model.inference(subgraph_loader, device,
progress_bar=True)
rename_profile_file(
model_name, dataset_name, str(batch_size),
str(layers), str(hidden_channels),
str(subgraph_loader.num_neighbors))