HeteroEdgeSampler

[qoffee]

In the examples I've seen how to sample vertices, but how to sample edges? (via HGTLoader)
For example, if we solve the problem "edge_prediction".
I am trying to adapt the following example(https://github.com/mims-harvard/graphml-tutorials/blob/master/05-graph-attention/heterogenenous-graph-attention.ipynb):
GAE(Encode(via HeteroConv), Decoder(via DistMult in GAE))

In other words with examples in PyG -> i want to adopt https://github.com/pyg-team/pytorch_geometric/blob/74245f3a680c1f6fd1944623e47d9e677b43e827/examples/rgcn_link_pred.py with HeteroConv(for several types of nodes and edges and use HGTLoader(for edges), because i have a big graph).

The question can be broken down into 2 sub-questions:

1. How to use HeteroConv in a link prediction task?
2. How do I use the HGTLoader here?

Thanks.

[rusty1s]

Adding better sampling support for link prediction is definitely something that I want to add. Currently, you can get around it by overriding the sampler class, i.e., as we do in the OGB example.

For applying HeteroConv, you should be able to simply replace the RGCNConv calls with HeteroConv ones, which now operate on dictionaries of node and edge attributes.

[qoffee]

Sorry, i have one more question. When I define my graph, like this:

data_h = HeteroData()
data_h['node_type_1'].x = data.x[node_type==1]
data_h['node_type_2'].x = data.x[node_type==2]
data_h['node_type_3'].x = data.x[node_type==3]

data_h['node_type_1', 'edge_t_1', 'node_type_1'].edge_index = data.edge_index[:, data.edge_type==1]
data_h['node_type_2', 'edge_t_2', 'node_type_2'].edge_index = data.edge_index[:, data.edge_type==2]
data_h['node_type_2', 'edge_t_3', 'node_type_1'].edge_index = data.edge_index[:, data.edge_type==3]
data_h['node_type_1', 'edge_t_4', 'node_type_2'].edge_index = data.edge_index[:, data.edge_type==4]
data_h['node_type_1', 'edge_t_5', 'node_type_3'].edge_index = data.edge_index[:, (data.edge_type==5)]
data_h['node_type_3', 'edge_t_6', 'node_type_1'].edge_index = data.edge_index[:, (data.edge_type==6)]

Is it correct to use edge_index from one general pytorch.data that is one large graph? Or for each pair of entities, i need to define its own data with its own edge_index? The question arose because I am getting an error and I have edge_index from a common data:
IndexError: index out of range in self

[rusty1s]

No that is not correct. Edge indices in HeteroData hold local node indices from different types. As such, you need to decrement the indices such that edge indices in edge_type (node_type_x, edge_type, node_type_y) obey to the following conditions:

edge_index[0].max() < data['node_type_x'].num_nodes
edge_index[1].max() < data['node_type_y'].num_nodes

In case data.x holds nodes for each node type in coalesced chunks, you can simply decrement your local edge indices, e.g.:

data_h['node_type_2', 'edge_t_3', 'node_type_1'].edge_index = data.edge_index[:, data.edge_type==3] - torch.tensor([[data['node_type_2'].num_nodes, [0]])
data_h['node_type_1', 'edge_t_5', 'node_type_3'].edge_index = data.edge_index[:, (data.edge_type==5)] - torch.tensor([[0], [data['node_type_1'].num_nodes + data['node_type_2'].num_nodes]])

We have plans to support this feature in the future out-of-the-box :)