Deep Graph Contrastive Representation Learning (GRACE)

class hivegraph.contrastive.grace.GRACE(num_features: int, hidden: int, num_layers: int, drop_edge_rate_1: float, drop_edge_rate_2: float, drop_feature_rate_1: float, drop_feature_rate_2: float, activation: ~typing.Callable = <function relu>, base_model: ~torch.nn.modules.module.Module = <class 'torch_geometric.nn.conv.gcn_conv.GCNConv'>, projection_dim: int = 128, tau: float | None = 0.5, model_name: str = 'GRACE', **kwargs)

Bases: Module

Implementation of deep GRAph Contrastive rEpresentation learning (GRACE).

References:

• https://arxiv.org/abs/2006.04131v2

• https://github.com/CRIPAC-DIG/GRACE

• https://wandb.ai/sauravmaheshkar/GRACE/reports/Graph-Contrastive-Learning–Vmlldzo2MzYxODEy

train_step(x: Tensor, edge_index: Tensor) → Tensor

Perform a single training step.

Parameters:

• x (torch.Tensor) – Node features.

• edge_index (torch.Tensor) – Edge indices.

Returns:

Loss.

Return type:

float

semi_loss(z1: Tensor, z2: Tensor) → Tensor

Compute the “semi_loss” between two given views.

Space Complexity: \(O(N^2)\)

\[l(u_i, v_i) = \log \frac{e^{\theta (u_i, v_i) / \tau}}{e^{\theta (u_i, v_i) / \tau} + \color{blue}{\sum_{k=1}^{N} \mathbb{1}_{[k \neq 1]} e^{\theta(u_i, v_k)/ \tau}} + \color{green}{\sum_{k=1}^{N}\mathbb{1}_{[k \neq 1]} e^{\theta(u_i, u_k)/ \tau}} }\]

• The equation in blue represents the loss between the inter-view negative pairs.

• The equation in green represents the loss between the intra-view negative pairs.

Parameters:

• z1 (torch.Tensor) – First set of views.

• z2 (torch.Tensor) – Second set of views.

Returns:

“semi_loss” between the two sets of views.

Return type:

torch.Tensor

batched_semi_loss(z1: Tensor, z2: Tensor, batch_size: int) → Tensor

Calculate the “semi_loss” between a batch of views

Space Complexity: \(O(BN)\)

Parameters:

• z1 (torch.Tensor) – First batch of views.

• z2 (torch.Tensor) – Second batch of views.

• batch_size (int) – Batch size.

Returns:

“semi_loss” between the two batches.

Return type:

torch.Tensor

loss(z1: Tensor, z2: Tensor, mean: bool | None = True, batch_size: int = 0) → Tensor

Compute the overall loss for all positive pairs.

Eqn(2) from the paper.

\[\mathcal{J} = \frac{1}{2N} \displaystyle \sum_{i=1}^{N} \left ( l(u_i, v_i) + l(v_i, u_i) \right)\]

References:

• https://arxiv.org/abs/2006.04131v2

Parameters:

• z1 (torch.Tensor) – First set of views.

• z2 (torch.Tensor) – Second set of views.

• mean (bool, optional) – If True, return the mean loss. Defaults to True.

• batch_size (int, optional) – Batch size. Defaults to 0.

Returns:

Overall loss.

Return type:

torch.Tensor

forward(x: Tensor, edge_index: Tensor) → Tensor

Compute a forward pass through the encoder module.

Parameters:

• x (torch.Tensor) – Node features.

• edge_index (torch.Tensor) – Edge indices.

Returns:

Representations from the encoder module.

Return type:

torch.Tensor

compute_cosine_sim(z1: Tensor, z2: Tensor) → Tensor

Compute the cosine similarity between two sets of views.

Parameters:

• z1 (torch.Tensor) – First set of views.

• z2 (torch.Tensor) – Second set of views.

Returns:

Cosine similarity between the two sets of views.

Return type:

torch.Tensor

project(z: Tensor) → Tensor

Project the representations to a lower-dimensional space.

This has been shown to enchance the expression power of the critic, For details refer to the section 3.2.1

References:

• https://arxiv.org/abs/2006.04131v2

Parameters:

z (torch.Tensor) – Representations from the encoder module.

Returns:

Projected representations.

Return type:

torch.Tensor

normalize_with_temp(x: Tensor) → Tensor

Normalize the given tensor with the temperature.

Parameters:

x (torch.Tensor) – Tensor to be normalized.

Returns:

Normalized tensor.

Return type:

torch.Tensor

reset_parameters() → None

Reset the parameters of the model.