RandomLinkSplit in custom HeteroData

[andreasala98]

Hi, and tanks for maintaining this library very active!
I am facing some trouble with my Heterogeneous graph on which I want to perform link prediction (infer on new nodes).
My graph (which I query directly from Neo4j) has 4 different node types (only 1 has features) and 5 different edge types (only 1 has features). I would like to embed my graph, but I got stuck while using the RandomLinkSplit function, and I can't undersrand why it doesn't work.
To load the graph I assign manually the different node and edge types (most of the attributes dataframes are None because I have no features):

data = HeteroData()

data['Patient'].x = pat_x # [num_patients, num_features_patient]
data['GPtime'].x = gp_x # It's none because i have no features
data['Diagnosis'].x = dia_x # It's none because i have no features 
data['NoacNode'].x = noac_x # It's none because i have no features

data['GPtime', 'evolves', 'GPtime'].edge_index = edge_gp_gp_index 
data['GPtime', 'first_noac', 'Patient'].edge_index = edge_gp_pat_index 
data['GPtime', 'prescribes', 'Diagnosis'].edge_index = edge_gp_dia_index 
data['Patient', 'suffers_from', 'Diagnosis'].edge_index = edge_pat_dia_index
data['Patient', 'gp_orig', "NoacNode"].edge_index = edge_pat_noac_index 

data['GPtime', 'evolves', 'GPtime'].edge_attr = edge_gp_gp_label  #no edge features-> None
data['GPtime', 'first_noac', 'Patient'].edge_attr = edge_gp_pat_label #no edge features-> None
data['GPtime', 'prescribes', 'Diagnosis'].edge_attr = edge_gp_dia_label    #not none!
data['Patient', 'suffers_from', 'Diagnosis'].edge_attr = edge_pat_dia_label  #no edge features-> None
data['Patient', 'gp_orig', "NoacNode"].edge_attr = edge_pat_noac_label   #no edge features-> None

and then to perform train test split:

import torch_geometric.transforms as T
data = T.ToUndirected()(data)  #why is this needed, anyway?

transform = T.RandomLinkSplit(
    is_undirected=False,
    add_negative_train_samples=True,
    num_val=0.1,
    num_test=0.1,
    neg_sampling_ratio=1.0,
    edge_types=[('Patient', 'suffers_from', 'Diagnosis'),
                          ('Patient', 'gp_orig', 'NoacNode'),
                          ('GPtime', 'prescribes', 'Diagnosis'),
                          ('GPtime', 'first_noac', 'Patient'),
                          ('GPtime', 'evolves', 'GPtime')
                           ],
    rev_edge_types=[('Diagnosis', 'rev_suffers_from', 'Patient'),
                                ('NoacNode', 'rev_gp_orig', 'Patient'),
                                ('GPtime', 'rev_evolves', 'GPtime'),
                                ('Patient', 'rev_first_noac', 'GPtime'),
                                ('Diagnosis', 'rev_prescribes', 'GPtime')
                               ],
)

train_data, val_data, test_data = transform(data)

but I get the error

TypeError: unsupported operand type(s) for *: 'Tensor' and 'NoneType'

Is this because I am trying to include in the tt split some edges that don't have any features?

If I don't transform my graph to undirected and I dont use the reverse edges, I get the following error:

~\Anaconda3\lib\site-packages\torch_geometric\transforms\random_link_split.py in __init__(self, num_val, num_test, is_undirected, key, split_labels, add_negative_train_samples, neg_sampling_ratio, disjoint_train_ratio, edge_types, rev_edge_types)
    105 
    106         if isinstance(edge_types, list):
--> 107             assert isinstance(rev_edge_types, list)
    108             assert len(edge_types) == len(rev_edge_types)
    109 

AssertionError: 

Maybe there is something very stupid that I didnt understand about graphs, but I can't seem to get my head around this; I've been fighting with this function for over a week now...

Thanks a lot!

[rusty1s]

Can you share a simple example to reproduce? I don‘t think this is related to having no features.

[andreasala98]

Hi @rusty1s thanks for answering. Unfortunately, I cannot share my data. I performed the split manually without recurring to the built in function.
Anyway, after the split, when I try to run a generic model I always get this error: (btw I switched to a Node classification problem so the split is now on the nodes)

model = HeteroGNN(data.metadata(), hidden_channels=64, out_channels=4,
                  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.005, weight_decay=0.001)




---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_24536/1154820153.py in <module>
      5 
      6 with torch.no_grad():  # Initialize lazy modules.
----> 7     out = model(data.x_dict, data.edge_index_dict)
      8 
      9 optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=0.001)

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~\AppData\Local\Temp/ipykernel_24536/3551936193.py in forward(self, x_dict, edge_index_dict)
     15     def forward(self, x_dict, edge_index_dict):
     16         for conv in self.convs:
---> 17             x_dict = conv(x_dict, edge_index_dict)
     18             x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()}
     19         return self.lin(x_dict['Patient'])

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~\Anaconda3\lib\site-packages\torch_geometric\nn\conv\hetero_conv.py in forward(self, x_dict, edge_index_dict, *args_dict, **kwargs_dict)
    121 
    122             if src == dst:
--> 123                 out = conv(x_dict[src], edge_index, *args, **kwargs)
    124             else:
    125                 out = conv((x_dict[src], x_dict[dst]), edge_index, *args,

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~\Anaconda3\lib\site-packages\torch_geometric\nn\conv\sage_conv.py in forward(self, x, edge_index, size)
     78 
     79         # propagate_type: (x: OptPairTensor)
---> 80         out = self.propagate(edge_index, x=x, size=size)
     81         out = self.lin_l(out)
     82 

~\Anaconda3\lib\site-packages\torch_geometric\nn\conv\message_passing.py in propagate(self, edge_index, size, **kwargs)
    307                         kwargs[arg] = decomp_kwargs[arg][i]
    308 
--> 309                 coll_dict = self.__collect__(self.__user_args__, edge_index,
    310                                              size, kwargs)
    311 

~\Anaconda3\lib\site-packages\torch_geometric\nn\conv\message_passing.py in __collect__(self, args, edge_index, size, kwargs)
    195                     assert len(data) == 2
    196                     if isinstance(data[1 - dim], Tensor):
--> 197                         self.__set_size__(size, 1 - dim, data[1 - dim])
    198                     data = data[dim]
    199 

~\Anaconda3\lib\site-packages\torch_geometric\nn\conv\message_passing.py in __set_size__(self, size, dim, src)
    161         the_size = size[dim]
    162         if the_size is None:
--> 163             size[dim] = src.size(self.node_dim)
    164         elif the_size != src.size(self.node_dim):
    165             raise ValueError(

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got -2)

I've tried the same model on a built in dataset, and it works perfectly. (I hope the problem is not caused by the train test split I am doing manually lol). I can print my heterodata object and it gives

HeteroData(
  Patient={
    x=[3230, 2],
    y=[3230],
    train_mask=[3230],
    val_mask=[3230],
    test_mask=[3230]
  },
  GPtime={
    x=[2194],
    train_mask=[2194],
    val_mask=[2194],
    test_mask=[2194]
  },
  Diagnosis={
    x=[22],
    train_mask=[22],
    val_mask=[22],
    test_mask=[22]
  },
  (GPtime, evolves, GPtime)={ edge_index=[2, 1853] },
  (GPtime, first_noac, Patient)={ edge_index=[2, 3230] },
  (GPtime, prescribes, Diagnosis)={ edge_index=[2, 9577] },
  (Patient, suffers_from, Diagnosis)={ edge_index=[2, 11323] }
)

that looks exactly like a built in dataset. Can you spot the bug?

[rusty1s]

Please ensure that node features x are always two-dimensional, e.g., via x = x.view(x.size(0), -1).

[andreasala98]

It is now working! Thanks for the trick, you really saved me a lot of time!