Predict data after training

[treyshotz]

We are currently creating a model for link prediction on a heterogeneous graph. Our goal is to enter a small graph and make predictions on this input graph. We are currently having problems when we are loading the data.
Our model is the following.

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(data[node_type].x.shape[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):
        x_dict = {
            node_type: self.lin_dict[node_type](x).relu_()
            for node_type, x in x_dict.items()
        }

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

        return x_dict


class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)

    def forward(self, z_dict, edge_label_index):
        row, col = edge_label_index
        z = torch.cat([z_dict['entity'][row], z_dict['match'][col]], dim=-1)

        z = self.lin1(z).relu()
        z = self.lin2(z)
        return z.view(-1)


class Model(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
        super().__init__()
        self.encoder = HGT(hidden_channels=hidden_channels, out_channels=out_channels,
                           num_heads=num_heads, num_layers=num_layers)
        self.decoder = EdgeDecoder(hidden_channels)

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict)
        return self.decoder(z_dict, edge_label_index)

We are using a private dataset. For loading during testing, we are using the following;

transform = T.RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    neg_sampling_ratio=0.5,
    add_negative_train_samples=True,
    edge_types=edge_types,
    rev_edge_types=reverse_edge_types,
    is_undirected=True,
)

train_split, val_split, test_split = transform(data)

The training is done as following:

 pred = model(
        train_data.x_dict,
        train_data.edge_index_dict,
        train_data["entity", "match"].edge_label_index,
    )
    target = train_data["entity", "match"].edge_label
    loss = F.mse_loss(pred, target)

When we are loading the data we want to predict we are doing the following:

  pred_dataset = ds.Prediction()
  pred_dataset.process()
  pred_data = pred_dataset[0]

  for node_type in pred_data.node_types:
      pred_data[node_type].x = pred_data[node_type].x.float()
  pred_data[("xxx", "to", "yyy")].edge_index = pred_data[
      ("xxx", "to", "yyy")
  ].edge_index.long()

  edge_types = pred_data.edge_types

  # Add reverse edges for message passing
  pred_data = T.ToUndirected(merge=False)(pred_data)

  reverse_edge_types = [(x[2], "{}_{}".format('rev', x[1]), x[0]) for x in edge_types]

  pred = model(
       pred_split.x_dict,
       pred_split.edge_index_dict,
       pred_split["entity", "match"].edge_label_index,
   )

What we see is that there is no "edge_label_index" when we are not using the RandomLinkSplit. How could we predict then?

We tried using the RandomLinkSplit but got an error stating that we didn't have enough edges(?)

[rusty1s]

Yes, the edge_label_index and edge_label are generated by the RandomLinkSplit transform, and represent a subset of original edges used for training the model in a supervised fashion. If you want to use your model for prediction, edge_label_index denotes which edges you want to predict. As such, it is the task of the user to provide this information. For example, if you want to predict the existence of edges between node 0 and all other nodes, you can do:

src = torch.zeros(data.num_nodes, dtype=torch.long)
dst = torch.arange(data.num_nodes, dtype=torch.long)
edge_label_index = torch.stack([src, dst], dim=0)