dig.sslgraph.method

Pre-implemented Contrastive Methods

Contrastive method classes under dig.sslgraph.method, or alias dig.sslgraph.method.contrastive.model.

Contrastive

Base class for creating contrastive learning models for either graph-level or node-level tasks.

GRACE

Contrastive learning method proposed in the paper Deep Graph Contrastive Representation Learning.

GraphCL

Contrastive learning method proposed in the paper Graph Contrastive Learning with Augmentations.

InfoGraph

Contrastive learning method proposed in the paper InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization.

MVGRL

Contrastive learning method for graph-level tasks proposed in the paper Contrastive Multi-View Representation Learning on Graphs (MVGRL).

NodeMVGRL

Contrastive learning method for node-level tasks proposed in the paper Contrastive Multi-View Representation Learning on Graphs (MVGRL).

class Contrastive(objective, views_fn, graph_level=True, node_level=False, z_dim=None, z_n_dim=None, proj=None, proj_n=None, neg_by_crpt=False, tau=0.5, device=None, choice_model='last', model_path='models')[source]

Base class for creating contrastive learning models for either graph-level or node-level tasks.

Alias: dig.sslgraph.method.contrastive.model.Contrastive.

Parameters
  • objective (string, or callable) – The learning objective of contrastive model. If string, should be one of ‘NCE’ and ‘JSE’. If callable, should take lists of representations as inputs and returns loss Tensor (see dig.sslgraph.method.contrastive.objectives for examples).

  • views_fn (list of callable) – List of functions to generate views from given graphs.

  • graph_level (bool, optional) – Whether to include graph-level representation for contrast. (default: True)

  • node_level (bool, optional) – Whether to include node-level representation for contrast. (default: False)

  • z_dim (int, optional) – The dimension of graph-level representations. Required if graph_level = True. (default: None)

  • z_dim – The dimension of node-level representations. Required if node_level = True. (default: None)

  • proj (string, or Module, optional) – Projection head for graph-level representation. If string, should be one of "linear" or "MLP". Required if graph_level = True. (default: None)

  • proj_n (string, or Module, optional) – Projection head for node-level representations. If string, should be one of "linear" or "MLP". Required if node_level = True. (default: None)

  • neg_by_crpt (bool, optional) – The mode to obtain negative samples in JSE. If True, obtain negative samples by performing corruption. Otherwise, consider pairs of different graph samples as negative pairs. Only used when objective = "JSE". (default: False)

  • tau (int) – The tempurature parameter in InfoNCE (NT-XENT) loss. Only used when objective = "NCE". (default: 0.5)

  • device (int, or torch.device, optional) – The device to perform computation.

  • choice_model (string, optional) – Whether to yield model with best training loss or at the last epoch. (default: last)

  • model_path (string, optinal) – The directory to restore the saved model. (default: models)

train(encoder, data_loader, optimizer, epochs, per_epoch_out=False)[source]

Perform contrastive training and yield trained encoders per epoch or after the last epoch.

Parameters
  • encoder (Module, or list of Module) – A graph encoder shared by all views or a list of graph encoders dedicated for each view. If node_level = False, the encoder should return tensor of shape [n_graphs, z_dim]. Otherwise, return tuple of shape ([n_graphs, z_dim], [n_nodes, z_n_dim]) representing graph-level and node-level embeddings.

  • dataloader (Dataloader) – Dataloader for unsupervised learning or pretraining.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in encoder(s).

  • epochs (int) – Number of total training epochs.

  • per_epoch_out (bool) – If True, yield trained encoders per epoch. Otherwise, only yield the final encoder at the last epoch. (default: False)

Return type

generator.

class GRACE(dim, dropE_rate_1, dropE_rate_2, maskN_rate_1, maskN_rate_2, **kwargs)[source]

Contrastive learning method proposed in the paper Deep Graph Contrastive Representation Learning. You can refer to the benchmark code for an example of usage.

Alias: dig.sslgraph.method.contrastive.model.GRACE.

Parameters
  • dim (int) – The embedding dimension.

  • dropE_rate_1 (float) – The ratio of the edge dropping augmentation for view 1. A number between [0,1).

  • dropE_rate_2 (float) – The ratio of the edge dropping augmentation for view 1. A number between [0,1).

  • maskN_rate_1 (float) – The ratio of the node masking augmentation for view 2. A number between [0,1).

  • maskN_rate_2 (float) – The ratio of the node masking augmentation for view 2. A number between [0,1).

  • **kwargs (optinal) – Additional arguments of dig.sslgraph.method.Contrastive.

train(encoders, data_loader, optimizer, epochs, per_epoch_out=False)[source]

Perform contrastive training and yield trained encoders per epoch or after the last epoch.

Parameters
  • encoder (Module, or list of Module) – A graph encoder shared by all views or a list of graph encoders dedicated for each view. If node_level = False, the encoder should return tensor of shape [n_graphs, z_dim]. Otherwise, return tuple of shape ([n_graphs, z_dim], [n_nodes, z_n_dim]) representing graph-level and node-level embeddings.

  • dataloader (Dataloader) – Dataloader for unsupervised learning or pretraining.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in encoder(s).

  • epochs (int) – Number of total training epochs.

  • per_epoch_out (bool) – If True, yield trained encoders per epoch. Otherwise, only yield the final encoder at the last epoch. (default: False)

Return type

generator.

class GraphCL(dim, aug_1=None, aug_2=None, aug_ratio=0.2, **kwargs)[source]

Contrastive learning method proposed in the paper Graph Contrastive Learning with Augmentations. You can refer to the benchmark code for an example of usage.

Alias: dig.sslgraph.method.contrastive.model.GraphCL.

Parameters
  • dim (int) – The embedding dimension.

  • aug1 (sting, optinal) – Types of augmentation for the first view from ("dropN", "permE", "subgraph", "maskN", "random2", "random3", "random4"). (default: None)

  • aug2 (sting, optinal) – Types of augmentation for the second view from ("dropN", "permE", "subgraph", "maskN", "random2", "random3", "random4"). (default: None)

  • aug_ratio (float, optional) – The ratio of augmentations. A number between [0,1).

  • **kwargs (optinal) – Additional arguments of dig.sslgraph.method.Contrastive.

train(encoders, data_loader, optimizer, epochs, per_epoch_out=False)[source]

Perform contrastive training and yield trained encoders per epoch or after the last epoch.

Parameters
  • encoder (Module, or list of Module) – A graph encoder shared by all views or a list of graph encoders dedicated for each view. If node_level = False, the encoder should return tensor of shape [n_graphs, z_dim]. Otherwise, return tuple of shape ([n_graphs, z_dim], [n_nodes, z_n_dim]) representing graph-level and node-level embeddings.

  • dataloader (Dataloader) – Dataloader for unsupervised learning or pretraining.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in encoder(s).

  • epochs (int) – Number of total training epochs.

  • per_epoch_out (bool) – If True, yield trained encoders per epoch. Otherwise, only yield the final encoder at the last epoch. (default: False)

Return type

generator.

class InfoGraph(g_dim, n_dim, **kwargs)[source]

Contrastive learning method proposed in the paper InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. You can refer to the benchmark code for an example of usage.

Alias: dig.sslgraph.method.contrastive.model.InfoGraph.

Parameters
  • g_dim (int) – The embedding dimension for graph-level (global) representations.

  • n_dim (int) – The embedding dimension for node-level (local) representations. Typically, when jumping knowledge is included in the encoder, we have g_dim = n_layers * n_dim.

  • **kwargs (optinal) – Additional arguments of dig.sslgraph.method.Contrastive.

train(encoders, data_loader, optimizer, epochs, per_epoch_out=False)[source]

Perform contrastive training and yield trained encoders per epoch or after the last epoch.

Parameters
  • encoder (Module, or list of Module) – A graph encoder shared by all views or a list of graph encoders dedicated for each view. If node_level = False, the encoder should return tensor of shape [n_graphs, z_dim]. Otherwise, return tuple of shape ([n_graphs, z_dim], [n_nodes, z_n_dim]) representing graph-level and node-level embeddings.

  • dataloader (Dataloader) – Dataloader for unsupervised learning or pretraining.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in encoder(s).

  • epochs (int) – Number of total training epochs.

  • per_epoch_out (bool) – If True, yield trained encoders per epoch. Otherwise, only yield the final encoder at the last epoch. (default: False)

Return type

generator.

class MVGRL(g_dim, n_dim, diffusion_type='ppr', alpha=0.2, t=5, **kwargs)[source]

Contrastive learning method for graph-level tasks proposed in the paper Contrastive Multi-View Representation Learning on Graphs (MVGRL). You can refer to the benchmark code for an example of usage.

Alias: dig.sslgraph.method.contrastive.model.MVGRL.

Parameters
  • g_dim (int) – The embedding dimension for graph-level (global) representations.

  • n_dim (int) – The embedding dimension for node-level (local) representations. Typically, when jumping knowledge is included in the encoder, we have g_dim = n_layers * n_dim.

  • diffusion_type (string, optional) – Diffusion instantiation mode with two options: "ppr": Personalized PageRank. "heat": heat kernel. (default: "ppr")

  • alpha (float, optional) – Teleport probability in a random walk. (default: 0.2)

  • t (int, optinal) – Diffusion time. (default: 5)

train(encoders, data_loader, optimizer, epochs, per_epoch_out=False)[source]

Perform contrastive training and yield trained encoders per epoch or after the last epoch.

Parameters
  • encoder (Module, or list of Module) – A graph encoder shared by all views or a list of graph encoders dedicated for each view. If node_level = False, the encoder should return tensor of shape [n_graphs, z_dim]. Otherwise, return tuple of shape ([n_graphs, z_dim], [n_nodes, z_n_dim]) representing graph-level and node-level embeddings.

  • dataloader (Dataloader) – Dataloader for unsupervised learning or pretraining.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in encoder(s).

  • epochs (int) – Number of total training epochs.

  • per_epoch_out (bool) – If True, yield trained encoders per epoch. Otherwise, only yield the final encoder at the last epoch. (default: False)

Return type

generator.

class NodeMVGRL(z_dim, z_n_dim, diffusion_type='ppr', alpha=0.2, t=5, batch_size=2, num_nodes=2000, **kwargs)[source]

Contrastive learning method for node-level tasks proposed in the paper Contrastive Multi-View Representation Learning on Graphs (MVGRL).

Alias: dig.sslgraph.method.contrastive.model.NodeMVGRL.

Parameters
  • z_g_dim (int) – The embedding dimension for graph-level (global) representations.

  • z_n_dim (int) – The embedding dimension for node-level (local) representations. Typically, when jumping knowledge is included in the encoder, we have z_g_dim = n_layers * z_n_dim.

  • diffusion_type (string, optional) – Diffusion instantiation mode with two options: "ppr": Personalized PageRank. "heat": heat kernel. (default: "ppr")

  • alpha (float, optional) – Teleport probability in a random walk. (default: 0.2)

  • t (int, optinal) – Diffusion time. (default: 5)

  • batch_size (int, optinal) – Number of subgraph samples in each minibatch. (default: 2)

  • num_nodes (int, optinal) – Number of nodes sampled in each subgraph. (default: 2000)

train(encoders, data_loader, optimizer, epochs, per_epoch_out=False)[source]

Perform contrastive training and yield trained encoders per epoch or after the last epoch.

Parameters
  • encoder (Module, or list of Module) – A graph encoder shared by all views or a list of graph encoders dedicated for each view. If node_level = False, the encoder should return tensor of shape [n_graphs, z_dim]. Otherwise, return tuple of shape ([n_graphs, z_dim], [n_nodes, z_n_dim]) representing graph-level and node-level embeddings.

  • dataloader (Dataloader) – Dataloader for unsupervised learning or pretraining.

  • optimizer (Optimizer) – Pytorch optimizer for trainable parameters in encoder(s).

  • epochs (int) – Number of total training epochs.

  • per_epoch_out (bool) – If True, yield trained encoders per epoch. Otherwise, only yield the final encoder at the last epoch. (default: False)

Return type

generator.

Contrastive Objectives

Contrastive objective functions under dig.sslgraph.method.contrastive.objectives.

JSE_loss(zs, zs_n=None, batch=None, sigma=None, neg_by_crpt=False, **kwargs)[source]

The Jensen-Shannon Estimator of Mutual Information used in contrastive learning. The implementation follows the paper Learning deep representations by mutual information estimation and maximization.

Note

The JSE loss implementation can produce negative values because a -2log2 shift is added to the computation of JSE, for the sake of consistency with other f-convergence losses.

Parameters
  • zs (list, optional) – List of tensors of shape [batch_size, z_dim].

  • zs_n (list, optional) – List of tensors of shape [nodes, z_dim].

  • batch (Tensor, optional) – Required when both zs and zs_n are given.

  • sigma (ndarray, optional) – A 2D-array of shape [n_views, n_views] with boolean values, indicating contrast between which two views are computed. Only required when number of views is greater than 2. If sigma[i][j] = True, JSE between \(view_i\) and \(view_j\) will be computed.

  • neg_by_crpt (bool, optional) – The mode to obtain negative samples in JSE. If True, obtain negative samples by performing corruption. Otherwise, consider pairs of different graph samples as negative pairs.

Return type

Tensor

NCE_loss(zs=None, zs_n=None, batch=None, sigma=None, **kwargs)[source]

The InfoNCE (NT-XENT) loss in contrastive learning.

Parameters
  • zs (list, optipnal) – List of tensors of shape [batch_size, z_dim].

  • zs_n (list, optional) – List of tensors of shape [nodes, z_dim].

  • batch (Tensor, optional) – Required when both zs and zs_n are given.

  • sigma (ndarray, optional) – A 2D-array of shape [n_views, n_views] with boolean values, indicating contrast between which two views are computed. Only required when number of views is greater than 2. If sigma[i][j] = True, infoNCE between \(view_i\) and \(view_j\) will be computed.

  • tau (int, optional) – The temperature used in NT-XENT.

Return type

Tensor

Transformations for Views Generation

Views generation classes under dig.sslgraph.method.contrastive.views_fn.

Diffusion

Diffusion on the given graph or batched graphs, used in MVGRL.

DiffusionWithSample

Diffusion with node sampling on the given graph for node-level datasets, used in MVGRL.

EdgePerturbation

Edge perturbation on the given graph or batched graphs.

NodeAttrMask

Node attribute masking on the given graph or batched graphs.

RWSample

Subgraph sampling based on random walk on the given graph or batched graphs.

RandomView

Generate views by random transformation (augmentation) on given batched graphs, where each graph in the batch is treated independently.

Sequential

Generate views by applying a sequence of transformations (augmentations) on given batched graphs.

UniformSample

Uniformly node dropping on the given graph or batched graphs.

class Diffusion(mode='ppr', alpha=0.2, t=5, add_self_loop=True)[source]

Diffusion on the given graph or batched graphs, used in MVGRL. Class objects callable via method views_fn().

Parameters
  • mode (string, optional) – Diffusion instantiation mode with two options: "ppr": Personalized PageRank; "heat": heat kernel. (default: "ppr")

  • alpha (float, optinal) – Teleport probability in a random walk. (default: 0.2)

  • t (float, optinal) – Diffusion time. (default: 5)

  • add_self_loop (bool, optional) – Set True to add self-loop to edge_index. (default: True)

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when Diffusion object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.

class DiffusionWithSample(sample_size=2000, batch_size=4, mode='ppr', alpha=0.2, t=5, epsilon=False, add_self_loop=True)[source]

Diffusion with node sampling on the given graph for node-level datasets, used in MVGRL. Class objects callable via method views_fn().

Parameters
  • sample_size (int, optional) – Number of nodes in the sampled subgraoh from a large graph. (default: 2000)

  • batch_size (int, optional) – Number of subgraphs to sample. (default: 4)

  • mode (string, optional) – Diffusion instantiation mode with two options: "ppr": Personalized PageRank; "heat": heat kernel; (default: "ppr")

  • alpha (float, optinal) – Teleport probability in a random walk. (default: 0.2)

  • t (float, optinal) – Diffusion time. (default: 5)

  • add_self_loop (bool, optional) – Set True to add self-loop to edge_index. (default: True)

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when DiffusionWithSample object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.

class EdgePerturbation(add=True, drop=False, ratio=0.1)[source]

Edge perturbation on the given graph or batched graphs. Class objects callable via method views_fn().

Parameters
  • add (bool, optional) – Set True if randomly add edges in a given graph. (default: True)

  • drop (bool, optional) – Set True if randomly drop edges in a given graph. (default: False)

  • ratio (float, optional) – Percentage of edges to add or drop. (default: 0.1)

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when EdgePerturbation object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.

class NodeAttrMask(mode='whole', mask_ratio=0.1, mask_mean=0.5, mask_std=0.5, return_mask=False)[source]

Node attribute masking on the given graph or batched graphs. Class objects callable via method views_fn().

Parameters
  • mode (string, optinal) – Masking mode with three options: "whole": mask all feature dimensions of the selected node with a Gaussian distribution; "partial": mask only selected feature dimensions with a Gaussian distribution; "onehot": mask all feature dimensions of the selected node with a one-hot vector. (default: "whole")

  • mask_ratio (float, optinal) – The ratio of node attributes to be masked. (default: 0.1)

  • mask_mean (float, optional) – Mean of the Gaussian distribution to generate masking values. (default: 0.5)

  • mask_std (float, optional) – Standard deviation of the distribution to generate masking values. Must be non-negative. (default: 0.5)

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when NodeAttrMask object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.

class RWSample(ratio=0.1, add_self_loop=False)[source]

Subgraph sampling based on random walk on the given graph or batched graphs. Class objects callable via method views_fn().

Parameters
  • ratio (float, optional) – Percentage of nodes to sample from the graph. (default: 0.1)

  • add_self_loop (bool, optional) – Set True to add self-loop to edge_index. (default: False)

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when RWSample object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.

class RandomView(candidates)[source]

Generate views by random transformation (augmentation) on given batched graphs, where each graph in the batch is treated independently. Class objects callable via method views_fn().

Parameters

candidates (list) – A list of callable view generation functions (classes).

__call__(data)[source]

Call self as a function.

views_fn(batch_data)[source]

Method to be called when RandomView object is called.

Parameters

batch_data (torch_geometric.data.Batch) – The input batched graphs.

Return type

torch_geometric.data.Batch.

class Sequential(fn_sequence)[source]

Generate views by applying a sequence of transformations (augmentations) on given batched graphs. Class objects callable via method views_fn().

Parameters

fn_sequence (list) – A list of callable view generation functions (classes).

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when Sequential object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.

class UniformSample(ratio=0.1)[source]

Uniformly node dropping on the given graph or batched graphs. Class objects callable via method views_fn().

Parameters

ratio (float, optinal) – Ratio of nodes to be dropped. (default: 0.1)

__call__(data)[source]

Call self as a function.

views_fn(data)[source]

Method to be called when UniformSample object is called.

Parameters

data (torch_geometric.data.Data) – The input graph or batched graphs.

Return type

torch_geometric.data.Data.