LinkPredictContrastiveDotDecoder

class graphstorm.model.LinkPredictContrastiveDotDecoder(in_dim)

Decoder for link prediction designed for contrastive loss by

using the dot product as the score function.

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

Parameters

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}.