LinkPredictContrastiveTransEDecoder

class graphstorm.model.LinkPredictContrastiveTransEDecoder(etypes, h_dim, gamma=12.0, norm='l2')

Decoder for link prediction designed for contrastive loss using the TransE 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.

etypes: list of tuples: The canonical edge types of the graph in the format of [(src_ntype1, etype1, dst_ntype1), …]
h_dim: int: The input dimension size. It is the dimension for both source and destination node embeddings.
gamma: float: The gamma value for model weight initialization. Default: 4.

forward(g, h, e_h=None)

Link prediction decoder forward function using the TransE

as the score function.

This computes the edge score on every edge type.

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