Feature dimension problem with explainer

[Exion35]

🐛 Describe the bug

I am currently performing a regression task on my graph with 12 vertices. My dataset consists of 1818 snapshots in the following format: Data(x=[12, 1, 336], edge_index=[2, 46], edge_attr=[46], y=[12, 48]). I have batched them with a size of 32, resulting in 56 elements in my dataloader.

I want to determine the important connections between my nodes, specifically identifying the nodes that contribute the most to the regression task. I'm attempting to use the following code (not fully reproducible as I can not give you the dataset, however the code is based on this repo and this notebook) to perform this explainability task :

import torch
from torch_geometric.explain import Explainer, PGExplainer
from torch_geometric_temporal.nn.attention import ASTGCN
import numpy as np
from tqdm import trange
from dataset import TemporalRegionalLoader # custom

DEVICE = torch.device('cpu') # cpu
batch_size = 32

adj_matrix_file = 'graph_representations/adjacency.txt' # 12x12 matrix
loader_train    = TemporalRegionalLoader('data', adj_matrix_file=adj_matrix_file)
dataset_train   = loader_train.get_dataset(num_timesteps_in=48*7, num_timesteps_out=48) # instance of the Static Graph Temporal Signal class (https://pytorch-geometric-temporal.readthedocs.io/en/latest/_modules/torch_geometric_temporal/signal/static_graph_temporal_signal.html#StaticGraphTemporalSignal)

train_input = np.array(dataset_train.features) # (1818, 12, 1, 336)
train_target = np.array(dataset_train.targets) # (1818, 12, 48)
train_x_tensor = torch.from_numpy(train_input).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
train_target_tensor = torch.from_numpy(train_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
train_dataset_new = torch.utils.data.TensorDataset(train_x_tensor, train_target_tensor)
train_loader = torch.utils.data.DataLoader(train_dataset_new, batch_size=batch_size, shuffle=False, drop_last=True)

model = ASTGCN(nb_block=3, in_channels=1, K=3, nb_chev_filter=64, nb_time_filter=64, time_strides=24, num_for_predict=48, len_input=48*7, num_of_vertices=12, bias=False).to(DEVICE)

explainer = Explainer(
    model=model,
    algorithm=PGExplainer(epochs=2, lr=1e-3),
    explanation_type='phenomenon',
    edge_mask_type='object',
    model_config=dict(
        mode='regression',
        task_level='graph',
        return_type='raw',
    ),
    # Include only the top 3 most important edges:
    threshold_config=dict(threshold_type='topk', value=3),
)

edge_index = torch.tensor(dataset_train.edge_index, dtype=torch.long)
for epoch in trange(2):
    for (x, target) in train_loader:
        loss = explainer.algorithm.train(
            epoch, model, x, edge_index, target=target)

However, I'm encountering the following error:

IndexError: The shape of the mask [46] at index 0 does not match the shape of the indexed tensor [552] at index 0.

Note that this model works just fine without calling the explainer. Thanks !

Environment

• PyG version: 2.3.0
• PyTorch version: 2.0.0+cpu
• OS: Windows
• Python version: 3.9.13
• How you installed PyTorch and PyG (conda, pip, source): pip

[rusty1s]

Can you share some dummy data I can work with to reproduce?

[Exion35]

Sure! Here is a csv file with dummy data : dummy_data.csv. The dataset is composed of a column date, node (values are in the range [1, 12]) and a feature column.

[rusty1s]

Do you have some guidance on how to load this in? It would be easier if you can just share x, edge_index and target that you input into the explainer.

[Exion35]

You should be able to reproduce the error with this code !

import pandas as pd
import numpy as np
from tqdm import trange
import torch
from torch_geometric.explain import Explainer, PGExplainer
from torch_geometric.utils import dense_to_sparse
from torch_geometric_temporal.signal import StaticGraphTemporalSignal
from torch_geometric_temporal.nn.attention import ASTGCN

PATH = 'data/dummy_data.csv'
DEVICE = torch.device('cpu') # cpu
batch_size = 32
A = torch.tensor([[0., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
                [1., 0., 1., 0., 1., 1., 1., 0., 0., 0., 0., 0.],
                [1., 1., 0., 1., 0., 0., 1., 1., 0., 0., 0., 0.],
                [1., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
                [0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
                [0., 1., 0., 0., 1., 0., 1., 0., 1., 0., 0., 0.],
                [0., 1., 1., 0., 0., 1., 0., 1., 1., 1., 0., 0.],
                [0., 0., 1., 1., 0., 0., 1., 0., 0., 1., 0., 0.],
                [0., 0., 0., 0., 0., 1., 1., 0., 0., 1., 1., 0.],
                [0., 0., 0., 0., 0., 0., 1., 1., 1., 0., 1., 1.],
                [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 0., 1.],
                [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 0.]], dtype=torch.float64)

df = pd.read_csv(PATH, index_col=0)
n_nodes = len(df.node.unique())

df_node = {}
for v in range(1, n_nodes+1):
    df_node[v] = df[df.node == v].reset_index(drop=True)

node_feats = np.array([[df_node[v].feature1.to_numpy()] for v in range(1, n_nodes+1)])
node_feats = torch.Tensor(node_feats).transpose(1, 0).unsqueeze(1)

edge_indices, values = dense_to_sparse(A)
edge_indices = edge_indices.numpy()
values = values.numpy()
indices = [(i, i + 8*48) for i in range(0, node_feats.shape[0] - 8*48 + 1, 48)]

features, targets = [], []
for i, j in indices:
    features.append((node_feats[i : i + 48*7, :].transpose(0, 2)).numpy())
    targets.append((node_feats[i + 48*7 : j, 0, :].transpose(0, 1)).numpy())

dataset_train = StaticGraphTemporalSignal(edge_indices, values, features, targets)

train_input = np.array(dataset_train.features)
train_target = np.array(dataset_train.targets) 
train_x_tensor = torch.from_numpy(train_input).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
train_target_tensor = torch.from_numpy(train_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
train_dataset_new = torch.utils.data.TensorDataset(train_x_tensor, train_target_tensor)
train_loader = torch.utils.data.DataLoader(train_dataset_new, batch_size=batch_size, shuffle=False, drop_last=True)

model = ASTGCN(nb_block=3, in_channels=1, K=3, nb_chev_filter=64, nb_time_filter=64, time_strides=24, num_for_predict=48, len_input=48*7, num_of_vertices=12, bias=False).to(DEVICE)

explainer = Explainer(
    model=model,
    algorithm=PGExplainer(epochs=2, lr=1e-3),
    explanation_type='phenomenon',
    edge_mask_type='object',
    model_config=dict(
        mode='regression',
        task_level='graph',
        return_type='raw',
    ),
    # Include only the top 3 most important edges:
    threshold_config=dict(threshold_type='topk', value=3),
)

edge_index = torch.tensor(dataset_train.edge_index, dtype=torch.long)
for epoch in trange(2):
    for (x, target) in train_loader:
        loss = explainer.algorithm.train(
            epoch, model, x, edge_index, target=target)

[rusty1s]

Thank you. This is helpful. On a first look, this seems to be a deeper issue which is non-trivial to fix, since the data format of PyG Temporal is quite different from what we expect in PyG and in our explainers (especially since we don't support dedicated batch sizes in node feature tensors in the explainer right now).