MLPEdgeDecoder

class graphstorm.model.MLPEdgeDecoder(h_dim, out_dim, multilabel, target_etype, num_hidden_layers=1, dropout=0, regression=False, num_ffn_layers=0, norm=None, use_bias=True)

Bases: GSEdgeDecoder

MLP-based decoder for edge prediction tasks.

Parameters

h_dim: int

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

out_dim: int

Output dimension size. If this decoder is for edge regression, the output dimension should be 1.

multilabel: bool

Whether this decoder is for multilabel edge classification.

target_etype: tuple of str

The target etype for prediction in the format of (src_ntype, etype, dst_ntype).

num_hidden_layers: int

Number of MLP layers. Default: 1.

dropout: float

Dropout rate. Default: 0.

regression: bool

Whether this decoder is for regression tasks. Default: False.

num_ffn_layers: int

Number of FFN layers added to the decoder. Default: 0

norm: str

Normalization methods. Not used, but reserved for complex MLPEdgeDecoder child class implementation. Default: None.

use_bias: bool

Whether the edge decoder uses a bias parameter. Default: True.

Changed in version 0.4.0: Add a new argument “use_bias” so users can control whether decoders have bias.

forward(g, h, e_h=None)

MLP-based edge decoder forward computation.

Parameters

g: DGLGraph

The graph of target edges.

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

out: Tensor

The prediction results.

predict(g, h, e_h=None)

MLP-based edge decoder predict computation.

Parameters

g: DGLGraph

The graph of target edges.

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

out: Tensor

The prediction results.

predict_proba(g, h, e_h=None)

MLP-based edge decoder prediction computation and return the normalized prediction results if this decoder is for edge classification.

Parameters

g: DGLGraph

The graph of target edges.

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

out: Tensor

The prediction results. If this decoder is for edge classification, return the normalized prediction results.

property in_dims

Return the input dimension size, which is given in class initialization.

property out_dims

Return the output dimension size. If this decoder is for edge regression, will return 1.