dig.sslgraph.method¶
Contents
Preimplemented Contrastive Methods¶
Contrastive method classes under dig.sslgraph.method
, or alias dig.sslgraph.method.contrastive.model
.
Base class for creating contrastive learning models for either graphlevel or nodelevel tasks. 

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

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

Contrastive learning method proposed in the paper InfoGraph: Unsupervised and Semisupervised GraphLevel Representation Learning via Mutual Information Maximization. 

Contrastive learning method for graphlevel tasks proposed in the paper Contrastive MultiView Representation Learning on Graphs (MVGRL). 

Contrastive learning method for nodelevel tasks proposed in the paper Contrastive MultiView Representation Learning on Graphs (MVGRL). 

Adaptive augmentation methods proposed in the paper Graph Contrastive Learning with Adaptive Augmentation. 
 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 graphlevel or nodelevel 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 graphlevel representation for contrast. (default:
True
)node_level (bool, optional) – Whether to include nodelevel representation for contrast. (default:
False
)z_dim (int, optional) – The dimension of graphlevel representations. Required if
graph_level
=True
. (default:None
)z_dim – The dimension of nodelevel representations. Required if
node_level
=True
. (default:None
)proj (string, or Module, optional) – Projection head for graphlevel representation. If string, should be one of
"linear"
or"MLP"
. Required ifgraph_level
=True
. (default:None
)proj_n (string, or Module, optional) – Projection head for nodelevel representations. If string, should be one of
"linear"
or"MLP"
. Required ifnode_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 (NTXENT) 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 thelast
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 graphlevel and nodelevel 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 graphlevel and nodelevel 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 graphlevel and nodelevel 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 Semisupervised GraphLevel 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 graphlevel (global) representations.
n_dim (int) – The embedding dimension for nodelevel (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 graphlevel and nodelevel 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 graphlevel tasks proposed in the paper Contrastive MultiView 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 graphlevel (global) representations.
n_dim (int) – The embedding dimension for nodelevel (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 graphlevel and nodelevel 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 nodelevel tasks proposed in the paper Contrastive MultiView Representation Learning on Graphs (MVGRL).
Alias:
dig.sslgraph.method.contrastive.model.
NodeMVGRL
. Parameters
z_g_dim (int) – The embedding dimension for graphlevel (global) representations.
z_n_dim (int) – The embedding dimension for nodelevel (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 graphlevel and nodelevel 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 pGRACE(dim: int, proj_n_dim: int, centrality_measure: str, prob_edge_1: float, prob_edge_2: float, prob_feature_1: float, prob_feature_2: float, tau: float = 0.1, dense: bool = False, p_tau: float = 0.7, **kwargs)[source]¶
Adaptive augmentation methods proposed in the paper Graph Contrastive Learning with Adaptive Augmentation. You can refer to the original implementation <https://github.com/CRIPACDIG/GCA> or the benchmark code for an example of usage.
Alias:
dig.sslgraph.method.contrastive.model.
pGRACE
. Parameters
dim (int) – The embedding dimension.
proj_n_dim (int) – The projection head dimension to use for computing loss
centrality_measure (str) – The metric to use for computing edge or node centrality. Supported values are degree, evc (EigenVector centrality) and pr (PageRank centrality).
prob_edge_1 (float) – The probability factor for calculating edgedrop probability
prob_edge_2 (float) – The probability factor for calculating edgedrop probability
prob_feature_1 (float) – The probability factor for calculating featuremasking probability
prob_feature_2 (float) – The probability factor for calculating featuremasking probability
tau (float, optional) – The temperature parameter used for contrastive objective.
dense (bool, optional) – Whether the node features are dense continuous features. Defaults to false.
p_tau (float, optional) – The upperbound probability for dropping edges or removing nodes.
**kwargs (optional) – 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 graphlevel and nodelevel 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 JensenShannon 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 fconvergence 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
andzs_n
are given.sigma (ndarray, optional) – A 2Darray 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 (NTXENT) 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
andzs_n
are given.sigma (ndarray, optional) – A 2Darray 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. Ifsigma[i][j]
=True
, infoNCE between \(view_i\) and \(view_j\) will be computed.tau (int, optional) – The temperature used in NTXENT.
 Return type
Tensor
Transformations for Views Generation¶
Views generation classes under dig.sslgraph.method.contrastive.views_fn
.
Proabilistic edge perturbation on the given graph or batched graphs. 

Proabilistic node attribute perturbation on the given graph or batched graphs. 

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

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

Edge perturbation on the given graph or batched graphs. 

Node attribute masking on the given graph or batched graphs. 

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

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

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

Uniformly node dropping on the given graph or batched graphs. 
 class AdaEdgePerturbation(centrality_measure: str, prob: float, threshold: float)[source]¶
Proabilistic edge perturbation on the given graph or batched graphs. Perturbations are adaptive i.e. critical edges have lower probability of being removed. Adaptations based on the paper Graph Contrastive Learning with Adaptive Augmentation <https://arxiv.org/abs/2010.14945> Refer to source implementation <https://github.com/CRIPACDIG/GCA> for more details. Class objects callable via method
views_fn()
. Parameters
centrality_measure (str) – Method for computing edge centrality. Set degree for degree centrality,
centrality (pr for PageRank centrality and evc for eigenvector) –
prob (float) – Probability factor used for calculating edgedrop probability
threshold (float) – Upperbound probability for dropping any edge, defaults to 0.7
 views_fn(data)[source]¶
Method to be called when
AdaEdgePerturbation
object is called. Parameters
data (
torch_geometric.data.Data
) – The input graph or batched graphs. Return type
 class AdaNodeAttrMask(centrality_measure: str, prob: float, threshold: float, dense: bool, return_mask: bool = False)[source]¶
Proabilistic node attribute perturbation on the given graph or batched graphs. Perturbations are adaptive i.e. critical nodes have lower probability of being perturbed. Adaptations based on the paper Graph Contrastive Learning with Adaptive Augmentation <https://arxiv.org/abs/2010.14945> Refer to source implementation <https://github.com/CRIPACDIG/GCA> for more details. Class objects callable via method
views_fn()
. Parameters
centrality_measure (str) – Method for computing node centrality. Set degree for degree centrality,
centrality (pr for PageRank centrality and evc for eigenvector) –
prob (float) – Probability factor used for calculating attributemasking probability
threshold (float, optional) – Upperbound probability for masking any attribute, defaults to 0.7
dense (bool, optional) – Indicates prescence of dense continuous features. Defaults to false
 views_fn(data)[source]¶
Method to be called when
AdaNodeAttrMask
object is called. Parameters
data (
torch_geometric.data.Data
) – The input graph or batched graphs. Return type
 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 selfloop to edge_index. (default:
True
)
 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
 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 nodelevel 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 selfloop to edge_index. (default:
True
)
 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
 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
 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
 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 onehot 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 nonnegative. (default:
0.5
)
 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
 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
 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
 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).
 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
 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).
 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
 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
)
 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