Graph Isomorphism Network (GIN)

class hivegraph.nn.gin.GIN(num_features: int, num_classes: int, num_layers: int, hidden: int, use_eps: bool = False, use_jump: bool = False, jump_mode: str = 'cat', dropout: float = 0.5, batchnorm: str = 'sequential', readout: str = 'mean', model_name: str = 'GIN')[source]

Bases: Module

Implementation of Graph Isomorphism Networks (GIN) from “How Powerful are Graph Neural Networks?” <https://arxiv.org/abs/1810.00826> by Keyulu Xu, Weihua Hu, Jure Leskovec, Stefanie Jegelka.

Parameters:
  • num_features (int) – Number of input features.

  • num_classes (int) – Number of output classes.

  • num_layers (int) – Number of GINConv layers.

  • hidden (int) – Number of hidden units.

  • use_eps (bool, optional) – If True, epsilon is a learnable parameter. Defaults to False.

  • use_jump (bool, optional) – If True, use JumpingKnowledge to aggregate representations from all layers. Defaults to False.

  • jump_mode (str, optional) – JumpingKnowledge aggregation mode. Must be one of ‘cat’, ‘max’, or ‘lstm’. Defaults to ‘cat’.

  • dropout (float, optional) – Dropout probability. Defaults to 0.5.

  • batchnorm (str, optional) – Batchnorm mode. Must be one of ‘first’, ‘last’, or ‘sequential’. Defaults to ‘last’.

  • readout (str, optional) – Readout function. Must be one of ‘mean’, ‘max’, or ‘sum’. Defaults to ‘mean’.

Supported batchnorm modes:

  • first: Batchnorm is applied to the input features of the first layer.

  • last: Batchnorm is applied to the output features of the last layer.

  • sequential: Batchnorm is applied to the output features of each layer.

Supported readout functions:

  • mean

  • max

  • sum

Supported JumpingKnowledge aggregation modes:

  • cat: Concatenate representations from all layers.

  • max: Take the maximum representation across all layers.

  • lstm: Use a LSTM to aggregate representations from all layers.

References

Raises:

AssertionError – If batchnorm, readout, or jump_mode is not supported.

reset_parameters() None[source]

Reset the parameters of the model.

forward(data: Data) Tensor[source]

Compute a forward pass through the model.

Parameters:

data (Data) – Data object.

Returns:

Logits for each class.

Return type:

torch.Tensor