Extracting Graph Embeddings from the train model

[sanchit-ahuja]

Hi, I wanted to know how can I extract the graph embeddings from the trained model for a custom graph input? Thanks!

[fedelopez77]

Hello,

You can do the following:

model = torch.load(path_to_saved_model)
embeddings = model["model"]["module.embeddings.embeds"]

Now embeddings has a tensor of shape (len(nodes), 2, n, n) where embeddings[:, 0] is the real part of the nxn symmetric matrix, and embeddings[:, 1] is the imaginary part of the matrix.

In case you train embeddings with models other than the Siegels spaces, such as Euclidean, hyperbolic, etc, then embeddings will be a tensor of shape (len(nodes), n).

[sanchit-ahuja]

Thanks for this. This works!