torch_geometric.nn.models.MaskLabel

class MaskLabel(num_classes: int, out_channels: int, method: str = 'add')[source]

Bases: Module

The label embedding and masking layer from the “Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification” paper.

Here, node labels y are merged to the initial node features x for a subset of their nodes according to mask.

Note

For an example of using MaskLabel, see examples/unimp_arxiv.py.

Parameters:
  • num_classes (int) – The number of classes.

  • out_channels (int) – Size of each output sample.

  • method (str, optional) – If set to "add", label embeddings are added to the input. If set to "concat", label embeddings are concatenated. In case method="add", then out_channels needs to be identical to the input dimensionality of node features. (default: "add")

forward(x: Tensor, y: Tensor, mask: Tensor) Tensor[source]
reset_parameters()[source]

Resets all learnable parameters of the module.

static ratio_mask(mask: Tensor, ratio: float)[source]

Modifies mask by setting ratio of True entries to False. Does not operate in-place.

Parameters:
  • mask (torch.Tensor) – The mask to re-mask.

  • ratio (float) – The ratio of entries to keep.