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])