torch_geometric.nn.models.GAE

class GAE(encoder: Module, decoder: Optional[Module] = None)

Bases: Module

The Graph Auto-Encoder model from the “Variational Graph Auto-Encoders” paper based on user-defined encoder and decoder models.

Parameters:

• encoder (torch.nn.Module) – The encoder module.

• decoder (torch.nn.Module, optional) – The decoder module. If set to None, will default to the torch_geometric.nn.models.InnerProductDecoder. (default: None)

forward(*args, **kwargs) → Tensor

Alias for encode().

reset_parameters()

Resets all learnable parameters of the module.

encode(*args, **kwargs) → Tensor

Runs the encoder and computes node-wise latent variables.

decode(*args, **kwargs) → Tensor

Runs the decoder and computes edge probabilities.

recon_loss(z: Tensor, pos_edge_index: Tensor, neg_edge_index: Optional[Tensor] = None) → Tensor

Given latent variables z, computes the binary cross entropy loss for positive edges pos_edge_index and negative sampled edges.

Parameters:

• z (torch.Tensor) – The latent space \(\mathbf{Z}\).

• pos_edge_index (torch.Tensor) – The positive edges to train against.

• neg_edge_index (torch.Tensor, optional) – The negative edges to train against. If not given, uses negative sampling to calculate negative edges. (default: None)

test(z: Tensor, pos_edge_index: Tensor, neg_edge_index: Tensor) → Tuple[Tensor, Tensor]

Given latent variables z, positive edges pos_edge_index and negative edges neg_edge_index, computes area under the ROC curve (AUC) and average precision (AP) scores.

Parameters:

• z (torch.Tensor) – The latent space \(\mathbf{Z}\).

• pos_edge_index (torch.Tensor) – The positive edges to evaluate against.

• neg_edge_index (torch.Tensor) – The negative edges to evaluate against.