• Search by Module
• Search by Words
• Search Projects

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• cogdl-master
  • cogdl
    • datasets
      • gtn_data.py
      • modelnet.py
      • gatne.py
      • __init__.py
      • README.md
      • han_data.py
      • matlab_matrix.py
      • pyg.py
    • layers
      • maggregator.py
      • se_layer.py
      • __init__.py
      • mixhop_layer.py
    • models
      • base_model.py
      • nn
        • asgcn.py
        • pyg_drgat.py
        • fastgcn.py
        • pyg_infomax.py
        • pyg_cheb.py
        • pyg_dgcnn.py
        • han.py
        • pyg_drgcn.py
        • gin.py
        • patchy_san.py
        • pyg_gat.py
        • gat.py
        • gcn.py
        • diffpool.py
        • pyg_unet.py
        • mixhop.py
        • __init__.py
        • sortpool.py
        • gtn.py
        • infograph.py
        • mlp.py
        • graphsage.py
        • pyg_gcn.py
      • __init__.py
      • README.md
      • emb
        • spectral.py
        • sdne.py
        • grarep.py
        • line.py
        • metapath2vec.py
        • hin2vec.py
        • dngr.py
        • netmf.py
        • pte.py
        • netsmf.py
        • prone.py
        • node2vec.py
        • gatne.py
        • deepwalk.py
        • __init__.py
        • dgk.py
        • hope.py
        • graph2vec.py
    • options.py
    • data
      • extract.py
      • data.py
      • download.py
      • dataloader.py
      • makedirs.py
      • batch.py
      • __init__.py
      • dataset.py
    • tasks
      • graph_classification.py
      • unsupervised_node_classification.py
      • multiplex_link_prediction.py
      • node_classification.py
      • link_prediction.py
      • base_task.py
      • multiplex_node_classification.py
      • unsupervised_graph_classification.py
      • node_classification_sampling.py
      • __init__.py
      • README.md
      • heterogeneous_node_classification.py
    • utils.py
  • examples
    • gcn.py
  • LICENSE
  • .readthedocs.yml
  • setup.py
  • .travis.yml
  • README.md
  • tests
    • datasets
      • test_pyg.py
    • tasks
      • test_heterogeneous_node_classification.py
      • test_node_classification.py
      • test_multiplex_node_classification.py
      • test_graph_classification.py
      • test_unsupervised_graph_classification.py
      • test_unsupervised_node_classification.py
      • test_multiplex_link_prediction.py
      • test_node_classification_sampling.py
      • test_link_prediction.py
  • scripts
    • display_data.py
    • download.py
    • train.py
    • parallel_train.py
  • README_CN.md
  • .gitignore
  • docs
    • Makefile
    • source
      • install.rst
      • datasets.rst
      • task
        • node_classification.rst
        • unsupervised_node_classification.rst
        • supervised_graph_classification.rst
        • unsupervised_graph_classification.rst
        • index.rst
      • layers.rst
      • data.rst
      • citing.rst
      • tutorial.rst
      • tasks.rst
      • models.nn.rst
      • models.emb.rst
      • utils.rst
      • options.rst
      • license.rst
      • modules.rst
      • index.rst
      • conf.py
      • models.rst
      • _static
    • requirements.txt

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter

from .. import BaseModel, register_model


class GraphConvolution(nn.Module):
    """
    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
    """

    def __init__(self, in_features, out_features, bias=True):
        super(GraphConvolution, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter("bias", None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.weight.size(1))
        self.weight.data.normal_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.normal_(-stdv, stdv)

    def forward(self, input, edge_index):
        adj = torch.sparse_coo_tensor(
            edge_index,
            torch.ones(edge_index.shape[1]).float(),
            (input.shape[0], input.shape[0]),
        ).to(input.device)
        support = torch.mm(input, self.weight)
        output = torch.spmm(adj, support)
        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return (
            self.__class__.__name__
            + " ("
            + str(self.in_features)
            + " -> "
            + str(self.out_features)
            + ")"
        )


@register_model("gcn")
class TKipfGCN(BaseModel):
    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        # fmt: off
        parser.add_argument("--num-features", type=int)
        parser.add_argument("--num-classes", type=int)
        parser.add_argument("--hidden-size", type=int, default=64)
        parser.add_argument("--dropout", type=float, default=0.5)
        # fmt: on

    @classmethod
    def build_model_from_args(cls, args):
        return cls(args.num_features, args.hidden_size, args.num_classes, args.dropout)

    def __init__(self, nfeat, nhid, nclass, dropout):
        super(TKipfGCN, self).__init__()

        self.gc1 = GraphConvolution(nfeat, nhid)
        self.gc2 = GraphConvolution(nhid, nclass)
        self.dropout = dropout
        # self.nonlinear = nn.SELU()

    def forward(self, x, adj):
        x = F.relu(self.gc1(x, adj))
        # h1 = x
        x = F.dropout(x, self.dropout, training=self.training)
        x = self.gc2(x, adj)

        # x = F.relu(x)
        # x = torch.sigmoid(x)
        # return x
        # h2 = x
        return F.log_softmax(x, dim=-1)
    
    def loss(self, data):
        return F.nll_loss(
            self.forward(data.x, data.edge_index)[data.train_mask],
            data.y[data.train_mask],
        )
    
    def predict(self, data):
        return self.forward(data.x, data.edge_index)