LinkPredictRotatEDecoder

class graphstorm.model.LinkPredictRotatEDecoder(etypes, h_dim, gamma=12.0)

Bases: LinkPredictMultiRelationLearnableDecoder

Decoder for link prediction using the RotatE as the score function.

Score function of RotateE measures the angular distance between head and tail elements. The angular distance is defined as:

\[d_r(h, t)=\|h\circ r-t\|\]

The RotatE score function is defined as:

\[gamma - \|h\circ r-t\|^2\]

where gamma is a margin.

For more details, please refer to https://arxiv.org/abs/1902.10197 or https://dglke.dgl.ai/doc/kg.html#rotatee.

Note: The relation embedding of RotatE has two parts, one for real numbers and one for complex numbers. Each has the dimension size as half of the input dimension size.

Parameters

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 initialization and score function. Default: 12.

New in version 0.4.0: The LinkPredictRotatEDecoder.

init_w_relation()

Initialize learnable relation embeddings.

An example:

def init_w_relation(self):
    self._w_relation = nn.Embedding(self.num_rels, self.h_dim)

    nn.init.uniform_(self._w_relation.weight, -1., 1.)

forward(g, h, e_h=None)

Link prediction decoder forward function using the RotatE

as the score function.

This computes the edge score on every edge type.

Parameters

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.

Returns

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

calc_test_scores(emb, pos_neg_tuple, neg_sample_type, device)

Compute scores for positive edges and negative edges.

Parameters

emb: dict of Tensor

Node embeddings in the format of {ntype: emb}.

pos_neg_tuple: dict of tuple

Positive and negative edges stored in a dict of tuple in the format of {(“src_ntype1”, “etype1”, “dst_ntype1” ): (pos_src_idx, neg_src_idx, pos_dst_idx, neg_dst_idx)}.

The pos_src_idx represents the postive source node indexes in the format of Torch.Tensor. The neg_src_idx represents the negative source node indexes in the format of Torch.Tensor. The pos_dst_idx represents the postive destination node indexes in the format of Torch.Tensor. The neg_dst_idx represents the negative destination node indexes in the format of Torch.Tensor.

We define positive and negative edges as:

The positive edges: (pos_src_idx, pos_dst_idx)
The negative edges: (pos_src_idx, neg_dst_idx) and (neg_src_idx, pos_dst_idx)

neg_sample_type: str

Describe how negative samples are sampled. There are two options:

Uniform: For each positive edge, we sample K negative edges.
Joint: For one batch of positive edges, we sample K negative edges.

device: th.device

Device used to compute scores.

Returns

scores: dict of tuple: Return a dictionary of edge type’s positive scores and negative scores in the format of {(src_ntype, etype, dst_ntype): (pos_scores, neg_scores)}.

property in_dims: Return the input dimension size, which is given in class initialization.

property out_dims: Return 1 for link prediction tasks.