how do I get inference results with node id and node embedding?

[Chanrom]

hi,

I am trying GAT example and I see that the model only returns loss, pred and labels, how can I get node id and node embedding returned?

    def forward(self, inputs):
        """Evaluate model, calculate loss, predictions and extract labels."""
        # fmt: off
        nodes, feat, mask, labels, edges, edges_value, adj_shape = inputs
        nodes = torch.squeeze(nodes)                # [N], N: num of nodes in subgraph
        feat = torch.squeeze(feat)                  # [N, F]
        mask = torch.squeeze(mask)                  # [N]
        labels = torch.squeeze(labels)              # [N]
        edges = torch.squeeze(edges)                # [X, 2], X: num of edges in subgraph
        edges_value = torch.squeeze(edges_value)    # [X]
        adj_shape = torch.squeeze(adj_shape)        # [2]
        # fmt: on

        sp_adj = torch.sparse_coo_tensor(edges, edges_value, adj_shape.tolist())
        h_1 = self.input_layer(feat, sp_adj)
        scores = self.out_layer(h_1, sp_adj)

        labels = labels.type(torch.int64)
        labels = labels[mask]  # [batch_size]
        scores = scores[mask]  # [batch_size]
        pred = scores.argmax(dim=1)
        loss = self.xent(scores, labels)
        return loss, pred, labels

[coledie]

What is the goal of returning node id and embedding? If you are looking to run inference, you can add an extra function to run the model with your desired outputs.

[Chanrom]

thanks @coledie . I see trainer.py has inference step to return embedding, I will take a reference. thanks!