Can train a Heterogeneous GNN without any issues, but then get a matmul dimensions issue when explaining it?

[jbmcs]

Firstly want to say thanks for making such a rich GNN repository, this is fantastic.

I've been trying to train a Heterogeneous GNN based on the tutorials, on the following graph of TV watchers and TV programs.

There are edges from TV programs to TV programs, and edges from TV programs to TV watchers. I have created a dummy data object with this code:

num_tv_watchers = 100 
num_tv_programs = 50 

 
tv_watcher_features = torch.randn(num_tv_watchers, 16) 
tv_program_features = torch.randn(num_tv_programs, 8) 

tv_watcher_labels = torch.tensor(np.random.randint(0, 2, num_tv_watchers))

edges_program_to_program = torch.randint(0, num_tv_programs, (2, 1000)) 

program_to_watcher_dim0 = []
program_to_watcher_dim1 = []
for n in range(num_tv_programs):
    for m in range(num_tv_watchers):
        program_to_watcher_dim0.append(n)
        program_to_watcher_dim1.append(m)
        
edges_program_to_watcher = torch.tensor([program_to_watcher_dim0, program_to_watcher_dim1])


hetero_graph = HeteroData()

hetero_graph['tv_watcher'].x = tv_watcher_features
hetero_graph['tv_program'].x = tv_program_features

train_idx = list(np.random.randint(0, num_tv_watchers, int(0.8 * num_tv_watchers)))
val_idx = [n for n in range(num_tv_watchers) if n not in train_idx]
train_mask = torch.tensor([True if n in train_idx else False for n in range(num_tv_watchers)])
val_mask = torch.tensor([True if n in val_idx else False for n in range(num_tv_watchers)])

hetero_graph['tv_watcher'].train_mask = train_mask
hetero_graph['tv_watcher'].val_mask = val_mask
hetero_graph['tv_watcher'].test_mask = val_mask

hetero_graph['tv_watcher'].y = tv_watcher_labels

hetero_graph['tv_program', 'related_to', 'tv_program'].edge_index = edges_program_to_program
hetero_graph['tv_program', 'watched_by', 'tv_watcher'].edge_index = edges_program_to_watcher

data = hetero_graph

So my hetero_graph object looks like this when I print it:

HeteroData(
tv_watcher={
x=[100, 16],
train_mask=[100],
val_mask=[100],
test_mask=[100],
y=[100]
},
tv_program={ x=[50, 8] },
(tv_program, related_to, tv_program)={ edge_index=[2, 1000] },
(tv_program, watched_by, tv_watcher)={ edge_index=[2, 5000] }
)

I train the model with the following code (I am wondering whether there are errors here, but I can't see the wood from the trees at the moment)

class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
        super().__init__()

        self.lin_dict = torch.nn.ModuleDict()
        for node_type in data.node_types:
            self.lin_dict[node_type] = Linear(-1, hidden_channels)

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HGTConv(hidden_channels, hidden_channels, data.metadata(),
                           num_heads, group='sum')
            self.convs.append(conv)

        self.lin = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        for node_type, x in x_dict.items():
            x_dict[node_type] = self.lin_dict[node_type](x).relu_()

        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)

        return self.lin(x_dict['tv_watcher'])

model = HGT(hidden_channels=16, out_channels=1,
            num_heads=2, num_layers=2)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data, model = data.to(device), model.to(device)

with torch.no_grad():  # Initialize lazy modules.
     out = model(data.x_dict, data.edge_index_dict)

optimizer = torch.optim.Adam(model.parameters(), lr=0.05, weight_decay=0.001)

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x_dict, data.edge_index_dict)
    mask = data['tv_watcher'].train_mask
    loss = F.binary_cross_entropy_with_logits(out[mask].squeeze(),
                                              data['tv_watcher'].y[mask].float())
    loss.backward()
    optimizer.step()
    return float(loss)

@torch.no_grad()
def test():
    model.eval()
    pred = torch.sigmoid(model(data.x_dict, data.edge_index_dict)).squeeze()
    pred = (pred > 0.5).float()
    accs = []
    for split in ['train_mask', 'val_mask', 'test_mask']:
        mask = data['tv_watcher'][split]
        acc = (pred[mask] == data['tv_watcher'].y[mask].float()).sum() / mask.sum()
        accs.append(float(acc))
    return accs

for epoch in range(1, 20):
    loss = train()
    train_acc, val_acc, test_acc = test()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, Test: {test_acc:.4f}')

As this is dummy data for now, the results are obviously rubbish, but the train loss does decrease and the model overfits.
I then try to explain the prediction for a given node:

explainer = Explainer(
    model=model, 
    algorithm=CaptumExplainer('IntegratedGradients'),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config = dict(
        mode='binary_classification',
        task_level='node', 
        return_type='probs',  # Model returns probabilities.
    ),
)

hetero_explanation = explainer(
    data.x_dict,
    data.edge_index_dict,
    index=torch.tensor([1]),

But I get the error:

**RuntimeError                              Traceback (most recent call last)
[<ipython-input-37-8bc8f819834e>](https://localhost:8080/#) in <cell line: 16>()
     14 )
     15 
---> 16 hetero_explanation = explainer(
     17     data.x_dict,
     18     data.edge_index_dict,

14 frames
[/usr/local/lib/python3.9/dist-packages/torch_geometric/nn/dense/linear.py](https://localhost:8080/#) in forward(self, x)
    130             x (torch.Tensor): The input features.
    131         """
--> 132         return F.linear(x, self.weight, self.bias)
    133 
    134     @torch.no_grad()

RuntimeError: mat1 and mat2 shapes cannot be multiplied (50x16 and 8x16)**

How can I fix this? Is it just an issue with my original model, or do I need to modify the input to explainer()?

[wsad1]

I think you need to update your foward function, because currently the feature dict gets updated in place

def forward(x_dict, edge_index_dict):
        out_dict = {}
        for node_type, x in x_dict.items():
            out_dict[node_type] = self.lin_dict[node_type](x).relu_()
        for conv in self.convs:
            out_dict = conv(out_dict, edge_index_dict)

        return self.lin(out_dict['tv_watcher'])