def test_chickenpox():
loader = ChickenpoxDatasetLoader()
dataset = loader.get_dataset()
for epoch in range(3):
for snapshot in 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, )
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
return loss
def validation_step(self, val_batch, batch_idx):
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])