LinkPredictContrastiveDotDecoder

class graphstorm.model.LinkPredictContrastiveDotDecoder(in_dim)

Decoder for link prediction designed for contrastive loss by using the dot product as the score function.

Note:

This class is specifically implemented for contrastive loss. But it could also be used by other pair-wise loss functions for link prediction tasks.

in_dim: int: The input dimension size. It is the dimension for both source and destination node embeddings.

forward(g, h, e_h=None)

Link prediction dot product decoder forward function.

This function computes the edge scores on all edge types of the input graph.

g: DGLGraph: The input graph.
h: dict of Tensor: The input node embeddings in the format of {ntype: emb}.
e_h: dict of Tensor: The input edge embeddings in the format of {(src_ntype, etype, dst_ntype): emb}. Not used, but reserved for future support of edge embeddings. Default: None.

scores: dict of Tensor: The scores for edges of all edge types in the input graph in the format of {(src_ntype, etype, dst_ntype): score}.