models.nn.pyg_diffpool

Module Contents

Classes

EntropyLoss

LinkPredLoss

GraphSAGE

GraphSAGE from “Inductive Representation Learning on Large Graphs”.

BatchedGraphSAGE

GraphSAGE with mini-batch

BatchedDiffPoolLayer

DIFFPOOL from paper `”Hierarchical Graph Representation Learning

BatchedDiffPool

DIFFPOOL layer with batch forward

DiffPool

DIFFPOOL from paper `Hierarchical Graph Representation Learning

Functions

toBatchedGraph(batch_adj, batch_feat, node_per_pool_graph)

class models.nn.pyg_diffpool.EntropyLoss[source]

Bases: torch.nn.Module

forward(self, adj, anext, s_l)[source]
class models.nn.pyg_diffpool.LinkPredLoss[source]

Bases: torch.nn.Module

forward(self, adj, anext, s_l)[source]
class models.nn.pyg_diffpool.GraphSAGE(in_feats, hidden_dim, out_feats, num_layers, dropout=0.5, normalize=False, concat=False, use_bn=False)[source]

Bases: torch.nn.Module

GraphSAGE from “Inductive Representation Learning on Large Graphs”.

..math::

h^{i+1}_{mathcal{N}(v)}=AGGREGATE_{k}(h_{u}^{k}) h^{k+1}_{v} = sigma(mathbf{W}^{k}·CONCAT(h_{v}^{k}, h_{mathcal{N}(v)}))

Args:

in_feats (int) : Size of each input sample. hidden_dim (int) : Size of hidden layer dimension. out_feats (int) : Size of each output sample. num_layers (int) : Number of GraphSAGE Layers. dropout (float, optional) : Size of dropout, default: 0.5. normalize (bool, optional) : Normalze features after each layer if True, default: True.

forward(self, x, edge_index, edge_weight=None)[source]
class models.nn.pyg_diffpool.BatchedGraphSAGE(in_feats, out_feats, use_bn=True, self_loop=True)[source]

Bases: torch.nn.Module

GraphSAGE with mini-batch

Args:

in_feats (int) : Size of each input sample. out_feats (int) : Size of each output sample. use_bn (bool) : Apply batch normalization if True, default: True. self_loop (bool) : Add self loop if True, default: True.

forward(self, x, adj)[source]
class models.nn.pyg_diffpool.BatchedDiffPoolLayer(in_feats, out_feats, assign_dim, batch_size, dropout=0.5, link_pred_loss=True, entropy_loss=True)[source]

Bases: torch.nn.Module

DIFFPOOL from paper “Hierarchical Graph Representation Learning with Differentiable Pooling”.

\[X^{(l+1)} = S^{l)}^T Z^{(l)} A^{(l+1)} = S^{(l)}^T A^{(l)} S^{(l)} Z^{(l)} = GNN_{l, embed}(A^{(l)}, X^{(l)}) S^{(l)} = softmax(GNN_{l,pool}(A^{(l)}, X^{(l)}))\]
in_featsint

Size of each input sample.

out_featsint

Size of each output sample.

assign_dimint

Size of next adjacency matrix.

batch_sizeint

Size of each mini-batch.

dropoutfloat, optional

Size of dropout, default: 0.5.

link_pred_lossbool, optional

Use link prediction loss if True, default: True.

forward(self, x, edge_index, batch, edge_weight=None)[source]
get_loss(self)[source]
class models.nn.pyg_diffpool.BatchedDiffPool(in_feats, next_size, emb_size, use_bn=True, self_loop=True, use_link_loss=False, use_entropy=True)[source]

Bases: torch.nn.Module

DIFFPOOL layer with batch forward

in_featsint

Size of each input sample.

next_sizeint

Size of next adjacency matrix.

emb_sizeint

Dimension of next node feature matrix.

use_bnbool, optional

Apply batch normalization if True, default: True.

self_loopbool, optional

Add self loop if True, default: True.

use_link_lossbool, optional

Use link prediction loss if True, default: True.

use_entropybool, optioinal

Use entropy prediction loss if True, default: True.

forward(self, x, adj)[source]
get_loss(self)[source]
models.nn.pyg_diffpool.toBatchedGraph(batch_adj, batch_feat, node_per_pool_graph)[source]
class models.nn.pyg_diffpool.DiffPool(in_feats, hidden_dim, embed_dim, num_classes, num_layers, num_pool_layers, assign_dim, pooling_ratio, batch_size, dropout=0.5, no_link_pred=True, concat=False, use_bn=False)[source]

Bases: models.BaseModel

DIFFPOOL from paper Hierarchical Graph Representation Learning with Differentiable Pooling.

in_featsint

Size of each input sample.

hidden_dimint

Size of hidden layer dimension of GNN.

embed_dimint

Size of embeded node feature, output size of GNN.

num_classesint

Number of target classes.

num_layersint

Number of GNN layers.

num_pool_layersint

Number of pooling.

assign_dimint

Embedding size after the first pooling.

pooling_ratiofloat

Size of each poolling ratio.

batch_sizeint

Size of each mini-batch.

dropoutfloat, optional

Size of dropout, default: 0.5.

no_link_predbool, optional

If True, use link prediction loss, default: True.

static add_args(parser)[source]

Add model-specific arguments to the parser.

classmethod build_model_from_args(cls, args)[source]

Build a new model instance.

classmethod split_dataset(cls, dataset, args)[source]
reset_parameters(self)[source]
after_pooling_forward(self, gnn_layers, adj, x, concat=False)[source]
forward(self, batch)[source]
loss(self, prediction, label)[source]