Dimension mismatch

[xuzhang0112]

Hello, I am using the to_hetero method to implement HGT for a subgraph classification problem and have encountered some dimension mismatch issues. However, the to_hetero method encapsulates many operations, preventing me from printing tensor dimensions within the GNN. When I print dimensions externally, my node dimensions and edge indices seem normal. How can I obtain more detailed error information? Or could you offer some insights on how to solve this problem?

here is a simple demo to reproduce my error:

import torch
import torch.nn as nn
from torch_geometric.nn import GATConv, Linear, to_hetero
from torch_geometric.data import HeteroData

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x

class Network(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.device = config['device']
        self.hidden_size, self.num_labels = 256, 684
        self.embed = nn.Embedding(26008, self.hidden_size)
        # self.embed.load_state_dict(torch.load('bert/embedding-pth'))
        self.dropout = nn.Dropout(p=0.1)
        self.homo_gat = GAT(256, 684)
        self.pooling = nn.AdaptiveMaxPool1d(1)
        self.node_types = config['node_types']
        self.edge_types = config['edge_types']
        self.graph_metadata = (self.node_types, self.edge_types)
        self.hetero_gat = to_hetero(self.homo_gat, self. graph_metadata)
        self.loss = nn.CrossEntropyLoss()
        
    def forward(self, input_list, graph, label_list):
        graph.to(self.device)
        x_dict, edge_index_dict = graph.x_dict, graph.edge_index_dict
        for k, v in x_dict.items():
            if v.nelement() != 0:
                _ = self.embed(v).permute(0,2,1)
                _ = self.pooling(_).squeeze(-1)
                x_dict[k] = self.dropout (_)
            else:
                x_dict[k] = torch.empty((0,256), dtype=torch.float).to(self.device)
        for k,v in x_dict.items():
            print(k, v.shape, v.dtype)
        for k,v in edge_index_dict.items():
            if v.nelement() != 0:
                print(k,v,v.dtype)
        x_dict = self.hetero_gat(x_dict, edge_index_dict)
        logits = x_dict['patient']
        loss = self.loss(logits, label_list[0].to(self.device))
        # return logits
        return [logits], [loss]
    
if __name__ == '__main__':
    node_types = ['patient', 'symptom', 'part', 'thing',
                    'time', 'reason', 'history',
                    'check1', 'checkup1', 'check2', 'checkup2',
                    'positive', 'positive2', 'attributes',
                    'change', "change2", 'reason2', 'procedure']
    edge_types = [(x,'_',y) for x in node_types
                             for y in node_types]
    config = {'device': 'cpu',
              'node_types': node_types,
              'edge_types': edge_types}
    model = Network(config).to(config['device'])

    data = HeteroData()
    for node_type in node_types:
        data[node_type].x = torch.empty((0,20), dtype=torch.long)
    for edge_type in edge_types:
        data[edge_type].edge_index = torch.empty((2,0), dtype=torch.long)
    
    data['patient'].x =  torch.randint(0,10000,(1,20), dtype=torch.long)
    data['symptom'].x =  torch.randint(0,10000,(3,20), dtype=torch.long)
    data['part'].x =  torch.randint(0,10000,(1,20), dtype=torch.long)
    data['time'].x =  torch.randint(0,10000,(2,20), dtype=torch.long)
    data['reason'].x =  torch.randint(0,10000,(1,20), dtype=torch.long)
    data['positive'].x =  torch.randint(0,10000,(1,20), dtype=torch.long)
    data['attributes'].x =  torch.randint(0,10000,(1,20), dtype=torch.long)
    data['change'].x =  torch.randint(0,10000,(2,20), dtype=torch.long)
    data['reason2'].x =  torch.randint(0,10000,(2,20), dtype=torch.long)
    data['patient','_','symptom'].edge_index = torch.tensor([[0,0,0], [0,1,2]], dtype=torch.long)
    data['patient','_','positive'].edge_index = torch.tensor([[0],[0]],dtype=torch.long)
    data['patient','_', 'change'].edge_index = torch.tensor([[0],[1]],dtype=torch.long)
    data['symptom','_', 'part'].edge_index = torch.tensor([[1,2],[0,0]],dtype=torch.long)
    data['symptom','_','time'].edge_index = torch.tensor([[0,1],[0,1]],dtype=torch.long)
    data['symptom','_','reason'].edge_index = torch.tensor([[1],[0]],dtype=torch.long)
    data['symptom','_','attributes'].edge_index = torch.tensor([[1],[0]],dtype=torch.long)
    data['symptom','_','change'].edge_index = torch.tensor([[2],[0]],dtype=torch.long)
    data['reason','_','time'].edge_index = torch.tensor([[0],[1]],dtype=torch.long)
    data['change','_','reason2'].edge_index = torch.tensor([[0,1],[0,1]],dtype=torch.long)
    
    out = model([], data, [torch.tensor([0],dtype=torch.long)])
    print(out)

and then you can see the error:

[rusty1s]

Closing this in favor of #9087.