Extract embeddings

[mustfkeskin]

I train graphsage algorithm. I want to extract node embeddings

model = GraphSAGE(
    data.num_node_features,
    hidden_channels=64,
    num_layers=2,
).to(device)

I wrote a simple embeder extraction class.
Are h_src and h_dst node embeddings?
data.edge_label_index[1] is this also the node_ids am I right?

class Embedder:
    def __init__(self, model_path):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = GraphSAGE(
                                data.num_node_features,
                                hidden_channels=64,
                                num_layers=2,
                            ).to(self.device)
        self.model.load_state_dict(torch.load(model_path))
        self.model.eval()

    @torch.no_grad()
    def get_node_embeddings(self, data):
        #x, edge_index = data
        h = self.model(data.x, data.edge_index) 
        h_src = h[data.edge_label_index[0]]
        h_dst = h[data.edge_label_index[1]]
            
        return data.edge_label_index[0], h_src

I want to speed up my prediction speed
This is my base data loader method
This method have generate positive and negative samples so it will be slower

train_loader = LinkNeighborLoader(
    data,
    batch_size=256,
    shuffle=True,
    neg_sampling_ratio=1.0,
    num_neighbors=[10, 10],
)

First I create pd dataframe which contains just nodes.
Because linkdataloader class have more than 20 million edge
But i have 1 million node
I dont want to predict each row
I just want to get the h_srcs and move forward
I fixed the target to 0, it doesn't matter which node it is connected to, only my first column will work.
Is my approach correct? is there an easier way to do this?

embedding_df = pd.DataFrame(nodes, columns = ["source"])
embedding_df["target"] = 0
embedding_df["features"] = node_features

edge_index = torch.from_numpy(df[["source", "target"]].to_numpy())
x = torch.from_numpy(np.array(embedding_df["vectors"].tolist(), dtype="float32"))

node_dataset = Dataset(x=x, edge_index=edge_index)
node_dataloader = DataLoader(node_dataset, 
                               batch_size=512,
                               shuffle=False)

[rusty1s]

Your Embedder class looks absolutely correct (although it should probably be a torch.nn.Module). Otherwise, I am not 100% sure on what's your concrete issue. You should be able to just iterate over the LinkNeighborLoader and collect embeddings. Alternatively, you can just iterate over each node once (via NeighborLoader), and compute the embeddings only once per node. This way, you avoid repeated computation of node embeddings in Embedder.

[mustfkeskin]

from torch_geometric.nn import GraphSAGE
class GraphSAGE(BasicGNN):
class BasicGNN(torch.nn.Module):

Yes graphsage derived from basicgnn. that class is also implemented from torch.nn.

I understood that the LinkNeighborLoader class generates the samples necessary to train the binary classification to see if there is an edge between the two nodes.

I used this example to train the graphsage algorithm unsupervised way

[rusty1s]

Usually, PyTorch recommends to implement this via hooks. Alternatively, you can write a wrapper around BasicGNN to achieve this, e.g.:

class MyGNN(torch.nn.Module):
    def __init__(self):
        self.gnn = GraphsAGE(in_channels, hidden_channels, num_layers)
        self.lin = Linear(hidden_channels, out_channels)
    
    def forward(self, x, edge_index):
        x = self.gnn(x, edge_index)
        return self.lin(x), x

[mustfkeskin]

The linear layer has weights and bias parameters. does this exactly mean embedding.

[rusty1s]

Can you clarify what you mean? For me, embedding would refer to the output of the gnn (in this case x), while lin will use that embedding to derive at a prediction (here lin is a final classifier for example).

[mustfkeskin]

linear layer change the value of embeddings isnt it

[rusty1s]

IMO, the linear layer takes the embedding and produces a prediction from it.