def test_discrete_train_test_split_dynamic():
snapshot_count = 250
n_count = 100
feature_count = 32
edge_indices, edge_weights, features = generate_signal(250, 100, 32)
feature = features[0]
targets = [
np.random.uniform(0, 10, (n_count, )) for _ in range(snapshot_count)
]
dataset = DynamicGraphStaticSignal(edge_indices, edge_weights, feature,
targets)
train_dataset, test_dataset = temporal_signal_split(dataset, 0.8)
for epoch in range(2):
for snapshot in test_dataset:
assert snapshot.edge_index.shape[0] == 2
assert snapshot.edge_index.shape[1] == snapshot.edge_attr.shape[0]
assert snapshot.x.shape == (100, 32)
assert snapshot.y.shape == (100, )
for epoch in range(2):
for snapshot in train_dataset:
assert snapshot.edge_index.shape[0] == 2
assert snapshot.edge_index.shape[1] == snapshot.edge_attr.shape[0]
assert snapshot.x.shape == (100, 32)
assert snapshot.y.shape == (100, )
def test_discrete_train_test_split_dynamic_batch():
snapshot_count = 250
node_count = 100
feature_count = 32
graph_count = 10
edge_indices, edge_weights, features, targets, batches = generate_signal(
snapshot_count, node_count, feature_count, graph_count)
feature = features[0]
dataset = DynamicGraphStaticSignalBatch(edge_indices, edge_weights,
feature, targets, batches)
train_dataset, test_dataset = temporal_signal_split(dataset, 0.8)
for epoch in range(2):
for snapshot in test_dataset:
assert snapshot.edge_index.shape[0] == 2
assert snapshot.edge_index.shape[1] == snapshot.edge_attr.shape[0]
assert snapshot.x.shape == (node_count * graph_count,
feature_count)
assert snapshot.y.shape == (node_count * graph_count, )
for epoch in range(2):
for snapshot in train_dataset:
assert snapshot.edge_index.shape[0] == 2
assert snapshot.edge_index.shape[1] == snapshot.edge_attr.shape[0]
assert snapshot.x.shape == (node_count * graph_count,
feature_count)
assert snapshot.y.shape == (node_count * graph_count, )
def test_train_test_split_dynamic_graph_static_signal():
snapshot_count = 250
n_count = 100
feature_count = 32
edge_indices, edge_weights, features, additional_features = generate_signal(
250, 100, 32, ["optional1", "optional2"])
targets = [
np.random.uniform(0, 10, (n_count, )) for _ in range(snapshot_count)
]
dataset = StaticGraphTemporalSignal(edge_indices[0], edge_weights[0],
features, targets,
**additional_features)
train_dataset, test_dataset = temporal_signal_split(dataset, 0.8)
for epoch in range(2):
for snapshot in test_dataset:
assert snapshot.edge_index.shape[0] == 2
assert snapshot.edge_index.shape[1] == snapshot.edge_attr.shape[0]
assert snapshot.x.shape == (100, 32)
assert snapshot.y.shape == (100, )
assert getattr(snapshot, "optional1").shape == (100, 32)
assert getattr(snapshot, "optional2").shape == (100, 32)
for epoch in range(2):
for snapshot in train_dataset:
assert snapshot.edge_index.shape[0] == 2
assert snapshot.edge_index.shape[1] == snapshot.edge_attr.shape[0]
assert snapshot.x.shape == (100, 32)
assert snapshot.y.shape == (100, )
assert getattr(snapshot, "optional1").shape == (100, 32)
assert getattr(snapshot, "optional2").shape == (100, 32)
def test_train_test_split_dynamic_hetero_graph_static_signal_batch():
snapshot_count = 250
node_count = 100
feature_count = 32
graph_count = 10
edge_index_dicts, edge_weight_dicts, feature_dicts, target_dicts, batch_dicts = generate_heterogeneous_signal(
snapshot_count, node_count, feature_count, graph_count)
dataset = DynamicHeteroGraphStaticSignalBatch(edge_index_dicts,
edge_weight_dicts,
feature_dicts[0],
target_dicts, batch_dicts)
train_dataset, test_dataset = temporal_signal_split(dataset, 0.8)
for _ in range(2):
for snapshot in test_dataset:
assert len(snapshot.node_types) == 2
assert snapshot.node_types[0] == 'author'
assert snapshot.node_types[1] == 'paper'
assert snapshot.node_stores[0]['x'].shape == (node_count *
graph_count,
feature_count)
assert snapshot.node_stores[1]['x'].shape == (node_count *
graph_count,
feature_count)
assert snapshot.node_stores[0]['y'].shape == (node_count *
graph_count, )
assert snapshot.node_stores[1]['y'].shape == (node_count *
graph_count, )
assert len(snapshot.edge_types) == 1
assert snapshot.edge_types[0] == ('author', 'writes', 'paper')
assert snapshot.edge_stores[0].edge_index.shape[0] == 2
assert snapshot.edge_stores[0].edge_index.shape[
1] == snapshot.edge_stores[0].edge_attr.shape[0]
for _ in range(2):
for snapshot in train_dataset:
assert len(snapshot.node_types) == 2
assert snapshot.node_types[0] == 'author'
assert snapshot.node_types[1] == 'paper'
assert snapshot.node_stores[0]['x'].shape == (node_count *
graph_count,
feature_count)
assert snapshot.node_stores[1]['x'].shape == (node_count *
graph_count,
feature_count)
assert snapshot.node_stores[0]['y'].shape == (node_count *
graph_count, )
assert snapshot.node_stores[1]['y'].shape == (node_count *
graph_count, )
assert len(snapshot.edge_types) == 1
assert snapshot.edge_types[0] == ('author', 'writes', 'paper')
assert snapshot.edge_stores[0].edge_index.shape[0] == 2
assert snapshot.edge_stores[0].edge_index.shape[
1] == snapshot.edge_stores[0].edge_attr.shape[0]
def test_train_test_split_static_hetero_graph_temporal_signal():
snapshot_count = 250
n_count = 100
feature_count = 32
edge_index_dicts, edge_weight_dicts, feature_dicts, target_dicts, additional_feature_dicts = generate_heterogeneous_signal(
snapshot_count, n_count, feature_count, "optional1", "optional2")
dataset = StaticHeteroGraphTemporalSignal(edge_index_dicts[0],
edge_weight_dicts[0],
feature_dicts, target_dicts,
**additional_feature_dicts)
train_dataset, test_dataset = temporal_signal_split(dataset, 0.8)
for epoch in range(2):
for snapshot in test_dataset:
assert len(snapshot.node_types) == 2
assert snapshot.node_types[0] == 'author'
assert snapshot.node_types[1] == 'paper'
assert snapshot.node_stores[0]['x'].shape == (n_count,
feature_count)
assert snapshot.node_stores[1]['x'].shape == (n_count,
feature_count)
assert snapshot.node_stores[0]['y'].shape == (n_count, )
assert snapshot.node_stores[1]['y'].shape == (n_count, )
assert len(snapshot.edge_types) == 1
assert snapshot.edge_types[0] == ('author', 'writes', 'paper')
assert snapshot.edge_stores[0].edge_index.shape[0] == 2
assert snapshot.edge_stores[0].edge_index.shape[
1] == snapshot.edge_stores[0].edge_attr.shape[0]
assert snapshot.node_stores[1]['optional1'].shape == (
n_count, feature_count)
assert snapshot.node_stores[1]['optional2'].shape == (
n_count, feature_count)
for epoch in range(2):
for snapshot in train_dataset:
assert len(snapshot.node_types) == 2
assert snapshot.node_types[0] == 'author'
assert snapshot.node_types[1] == 'paper'
assert snapshot.node_stores[0]['x'].shape == (n_count,
feature_count)
assert snapshot.node_stores[1]['x'].shape == (n_count,
feature_count)
assert snapshot.node_stores[0]['y'].shape == (n_count, )
assert snapshot.node_stores[1]['y'].shape == (n_count, )
assert len(snapshot.edge_types) == 1
assert snapshot.edge_types[0] == ('author', 'writes', 'paper')
assert snapshot.edge_stores[0].edge_index.shape[0] == 2
assert snapshot.edge_stores[0].edge_index.shape[
1] == snapshot.edge_stores[0].edge_attr.shape[0]
assert snapshot.node_stores[1]['optional1'].shape == (
n_count, feature_count)
assert snapshot.node_stores[1]['optional2'].shape == (
n_count, feature_count)
def test_discrete_train_test_split_static():
loader = ChickenpoxDatasetLoader()
dataset = loader.get_dataset()
train_dataset, test_dataset = temporal_signal_split(dataset, 0.8)
for epoch in range(2):
for snapshot in train_dataset:
assert snapshot.edge_index.shape == (2, 102)
assert snapshot.edge_attr.shape == (102, )
assert snapshot.x.shape == (20, 4)
assert snapshot.y.shape == (20, )
for epoch in range(2):
for snapshot in test_dataset:
assert snapshot.edge_index.shape == (2, 102)
assert snapshot.edge_attr.shape == (102, )
assert snapshot.x.shape == (20, 4)
assert snapshot.y.shape == (20, )
# Implementation to train a regressor on the Hungarian Chickenpox Cases dataset to predict weekly cases reported by the countries using a Recurrent Graph Convolutional Network (R-GCN)
# Import libraries
import torch
from tqdm import tqdm
import torch.nn.functional as F
from torch_geometric_temporal.nn.recurrent import DCRNN
from torch_geometric_temporal.dataset import ChickenpoxDatasetLoader
from torch_geometric_temporal.signal import temporal_signal_split
# Load dataset
loader = ChickenpoxDatasetLoader()
dataset = loader.get_dataset()
# Split dataset into train and test
train_dataset, test_dataset = temporal_signal_split(dataset, train_ratio=0.2)
# Define a Recurrent Graph Convolutional Network (R-GCN)
class RecurrentGCN(torch.nn.Module):
def __init__(self, node_features):
super(RecurrentGCN, self).__init__()
self.recurrent = DCRNN(node_features, 32, 1)
self.linear = torch.nn.Linear(32, 1)
def forward(self, x, edge_index, edge_weight):
h = self.recurrent(x, edge_index, edge_weight)
h = F.relu(h)
h = self.linear(h)
return h
x = val_batch.x
y = val_batch.y.view(-1, 1)
edge_index = val_batch.edge_index
h = self.recurrent(x, edge_index)
h = F.relu(h)
h = self.linear(h)
loss = F.mse_loss(h, y)
metrics = {'val_loss': loss}
self.log_dict(metrics)
return metrics
loader = ChickenpoxDatasetLoader()
dataset_loader = loader.get_dataset(lags=32)
train_loader, val_loader = temporal_signal_split(dataset_loader,
train_ratio=0.2)
model = LitDiffConvModel(node_features=32, filters=16)
early_stop_callback = EarlyStopping(monitor='val_loss',
min_delta=0.00,
patience=10,
verbose=False,
mode='max')
trainer = pl.Trainer(callbacks=[early_stop_callback])
trainer.fit(model, train_loader, val_loader)
