How to add a second edge index to use as a target in link prediction?

[davidireland3]

I've tried adding a new attribute that I would ideally like to behave the same as the edge_index attribute when used with the Data Loader.
At the moment I pass a tuple of the graph I am learning over, g1, with the graph whose edges I want to use as a target for link prediction, g2, to the data loader and do e.g. pos_ids = g2.edge_index and get the neg_ids from the negative batch sampler.
However, maintaining two separate graph objects seems untidy and it'd be ideal if there is a way to have the edge index g2 as an attribute in g1 that when passed through the Data Loader handles the relabelling of the nodes to be in line with those in the batch, as is done with the standard edge_index.

[Padarn]

it'd be ideal if there is a way to have the edge index g2 as an attribute in g1 that when passed through the Data Loader handles the relabelling of the nodes to be in line with those in the batch, as is done with the standard edge_index.

I didn't totally follow this, could you provide an example of what it might look like?

[davidireland3]

Here is a snippet of my code currently:

def train(self):
    data = [(d.to(self.device), d1.to(self.device)) for d, d1 in self.data]  # here is the 'untidiness', `d` is the graph I am learning over and then I am having to add `d1` as another graph object as a proxy to get the edge targets that I want
    N = int(0.8 * len(data))
    random.shuffle(data)
    train_data = DataLoader(data[:N], batch_size=self.batch_size, shuffle=True)
    val_data = DataLoader(data[N:], batch_size=len(data[N:]), shuffle=False)

    for epoch in range(num_epochs):
        print(f"Epoch {epoch + 1}")
        for count, batch in enumerate(train_data):
            batch, proxy_graph = batch
            pos_ids_for_edges = proxy_graph.edge_index  # this is the targets that come from the proxy graph
            neg_ids_for_edges = batched_negative_sampling(proxy_graph.edge_index, proxy_graph.batch, self.negative_samples)
            targets = torch.cat((self.one.repeat(pos_ids_for_edges.size(1)), self.zero.repeat(neg_ids_for_edges.size(1))), dim=0)  # this a target for cross entropy loss

            self.edge_optimiser.zero_grad()
            edge_preds = self.edge_model.forward(batch, pos_ids_for_edges, neg_ids_for_edges)  # this is just a GNN that takes pos edge ids, which come from the proxy graph edges, and neg edge ids which come from the sampler
            edge_pred_loss = self.ce_loss_fn(edge_preds, targets)
            self.edge_losses.append(edge_pred_loss.item())
            edge_pred_loss.backward()
            self.edge_optimisers.step()

To give a concrete example, say I have two graphs g1 and g2, which I put in a DataLoader with batch_size=2. If I stored the 'positive' pairs of edges that I want to predict for link prediction, I could store in a tensor e.g. g1.targets = torch.LongTensor([[...], [...]]) (similar for g2). But when the DataLoader does it's magic of updating the batch.edge_index to be relabelled, it doesn't relabel the batch.targets as I would want it to, so the node labels would be off, hence why I have to use the proxy graph so that the DataLoader can relabel them for me. I think I could probably code this myself to relabel the nodes in batch.targets, but I doubt it would be the most efficient, and in general I think it could be a cool feature for PyG to have.

[rusty1s]

If you want to have a separate attribute similar to edge_index, you may need a custom Data object and override the __cat_dim__ and __inc__ functionalities. This tutorial goes over it on more depth.

[davidireland3]

Thanks, I just took a look at this. At a quick glance, it looks like the pairs of graphs example could be what I need, just without the node feature for the proxy graph. I'll give this a go, thanks.