import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.module import Module
from torch.nn.parameter import Parameter
from .. import BaseModel, register_model
from cogdl.utils import spmm
class GraphConvolutionBS(Module):
"""
GCN Layer with BN, Self-loop and Res connection.
"""
def __init__(
self,
in_features,
out_features,
activation=lambda x: x,
withbn=True,
withloop=True,
bias=True,
res=False,
):
"""
Initial function.
:param in_features: the input feature dimension.
:param out_features: the output feature dimension.
:param activation: the activation function.
:param withbn: using batch normalization.
:param withloop: using self feature modeling.
:param bias: enable bias.
:param res: enable res connections.
"""
super(GraphConvolutionBS, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.sigma = activation
self.res = res
# Parameter setting.
self.weight = Parameter(torch.FloatTensor(in_features, out_features))
self.self_weight = Parameter(torch.FloatTensor(in_features, out_features)) if withloop else None
self.bn = torch.nn.BatchNorm1d(out_features) if withbn else None
self.bias = Parameter(torch.FloatTensor(out_features)) if bias else None
self.reset_parameters()
def reset_parameters(self):
stdv = 1.0 / math.sqrt(self.weight.size(1))
self.weight.data.uniform_(-stdv, stdv)
if self.self_weight is not None:
stdv = 1.0 / math.sqrt(self.self_weight.size(1))
self.self_weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform_(-stdv, stdv)
def forward(self, graph, x):
support = torch.mm(x, self.weight)
output = spmm(graph, support)
# Self-loop
output = output + torch.mm(x, self.self_weight) if self.self_weight is not None else output
output = output + self.bias if self.bias is not None else output
# BN
output = self.bn(output) if self.bn is not None else output
# Res
return self.sigma(output) + input if self.res else self.sigma(output)
def __repr__(self):
return self.__class__.__name__ + " (" + str(self.in_features) + " -> " + str(self.out_features) + ")"
class GraphBaseBlock(Module):
"""
The base block for Multi-layer GCN / ResGCN / Dense GCN
"""
def __init__(
self,
in_features,
out_features,
nbaselayer,
withbn=True,
withloop=True,
activation=F.relu,
dropout=True,
aggrmethod="concat",
dense=False,
):
"""
The base block for constructing DeepGCN model.
:param in_features: the input feature dimension.
:param out_features: the hidden feature dimension.
:param nbaselayer: the number of layers in the base block.
:param withbn: using batch normalization in graph convolution.
:param withloop: using self feature modeling in graph convolution.
:param activation: the activation function, default is ReLu.
:param dropout: the dropout ratio.
:param aggrmethod: the aggregation function for baseblock, can be "concat" and "add". For "resgcn", the default
is "add", for others the default is "concat".
:param dense: enable dense connection
"""
super(GraphBaseBlock, self).__init__()
self.in_features = in_features
self.hiddendim = out_features
self.nhiddenlayer = nbaselayer
self.activation = activation
self.aggrmethod = aggrmethod
self.dense = dense
self.dropout = dropout
self.withbn = withbn
self.withloop = withloop
self.hiddenlayers = nn.ModuleList()
self.__makehidden()
if self.aggrmethod == "concat" and dense is False:
self.out_features = in_features + out_features
elif self.aggrmethod == "concat" and dense is True:
self.out_features = in_features + out_features * nbaselayer
elif self.aggrmethod == "add":
if in_features != self.hiddendim:
raise RuntimeError("The dimension of in_features and hiddendim should be matched in add model.")
self.out_features = out_features
elif self.aggrmethod == "nores":
self.out_features = out_features
else:
raise NotImplementedError("The aggregation method only support 'concat','add' and 'nores'.")
def __makehidden(self):
# for i in xrange(self.nhiddenlayer):
for i in range(self.nhiddenlayer):
if i == 0:
layer = GraphConvolutionBS(
self.in_features,
self.hiddendim,
self.activation,
self.withbn,
self.withloop,
)
else:
layer = GraphConvolutionBS(
self.hiddendim,
self.hiddendim,
self.activation,
self.withbn,
self.withloop,
)
self.hiddenlayers.append(layer)
def _doconcat(self, x, subx):
if x is None:
return subx
if self.aggrmethod == "concat":
return torch.cat((x, subx), 1)
elif self.aggrmethod == "add":
return x + subx
elif self.aggrmethod == "nores":
return x
def forward(self, graph, x):
h = x
denseout = None
# Here out is the result in all levels.
for gc in self.hiddenlayers:
denseout = self._doconcat(denseout, h)
h = gc(graph, h)
h = F.dropout(h, self.dropout, training=self.training)
if not self.dense:
return self._doconcat(h, x)
return self._doconcat(h, denseout) if denseout is not None else h
def get_outdim(self):
return self.out_features
def __repr__(self):
return "%s %s (%d - [%d:%d] > %d)" % (
self.__class__.__name__,
self.aggrmethod,
self.in_features,
self.hiddendim,
self.nhiddenlayer,
self.out_features,
)
class MultiLayerGCNBlock(Module):
"""
Muti-Layer GCN with same hidden dimension.
"""
def __init__(
self,
in_features,
out_features,
nbaselayer,
withbn=True,
withloop=True,
activation=F.relu,
dropout=True,
aggrmethod=None,
dense=None,
):
"""
The multiple layer GCN block.
:param in_features: the input feature dimension.
:param out_features: the hidden feature dimension.
:param nbaselayer: the number of layers in the base block.
:param withbn: using batch normalization in graph convolution.
:param withloop: using self feature modeling in graph convolution.
:param activation: the activation function, default is ReLu.
:param dropout: the dropout ratio.
:param aggrmethod: not applied.
:param dense: not applied.
"""
super(MultiLayerGCNBlock, self).__init__()
self.model = GraphBaseBlock(
in_features=in_features,
out_features=out_features,
nbaselayer=nbaselayer,
withbn=withbn,
withloop=withloop,
activation=activation,
dropout=dropout,
dense=False,
aggrmethod="nores",
)
def forward(self, graph, x):
return self.model.forward(graph, x)
def get_outdim(self):
return self.model.get_outdim()
def __repr__(self):
return "%s %s (%d - [%d:%d] > %d)" % (
self.__class__.__name__,
self.aggrmethod,
self.model.in_features,
self.model.hiddendim,
self.model.nhiddenlayer,
self.model.out_features,
)
class ResGCNBlock(Module):
"""
The multiple layer GCN with residual connection block.
"""
def __init__(
self,
in_features,
out_features,
nbaselayer,
withbn=True,
withloop=True,
activation=F.relu,
dropout=True,
aggrmethod=None,
dense=None,
):
"""
The multiple layer GCN with residual connection block.
:param in_features: the input feature dimension.
:param out_features: the hidden feature dimension.
:param nbaselayer: the number of layers in the base block.
:param withbn: using batch normalization in graph convolution.
:param withloop: using self feature modeling in graph convolution.
:param activation: the activation function, default is ReLu.
:param dropout: the dropout ratio.
:param aggrmethod: not applied.
:param dense: not applied.
"""
super(ResGCNBlock, self).__init__()
self.model = GraphBaseBlock(
in_features=in_features,
out_features=out_features,
nbaselayer=nbaselayer,
withbn=withbn,
withloop=withloop,
activation=activation,
dropout=dropout,
dense=False,
aggrmethod="add",
)
def forward(self, graph, x):
return self.model.forward(graph, x)
def get_outdim(self):
return self.model.get_outdim()
def __repr__(self):
return "%s %s (%d - [%d:%d] > %d)" % (
self.__class__.__name__,
self.aggrmethod,
self.model.in_features,
self.model.hiddendim,
self.model.nhiddenlayer,
self.model.out_features,
)
class DenseGCNBlock(Module):
"""
The multiple layer GCN with dense connection block.
"""
def __init__(
self,
in_features,
out_features,
nbaselayer,
withbn=True,
withloop=True,
activation=F.relu,
dropout=True,
aggrmethod="concat",
dense=True,
):
"""
The multiple layer GCN with dense connection block.
:param in_features: the input feature dimension.
:param out_features: the hidden feature dimension.
:param nbaselayer: the number of layers in the base block.
:param withbn: using batch normalization in graph convolution.
:param withloop: using self feature modeling in graph convolution.
:param activation: the activation function, default is ReLu.
:param dropout: the dropout ratio.
:param aggrmethod: the aggregation function for the output. For denseblock, default is "concat".
:param dense: default is True, cannot be changed.
"""
super(DenseGCNBlock, self).__init__()
self.model = GraphBaseBlock(
in_features=in_features,
out_features=out_features,
nbaselayer=nbaselayer,
withbn=withbn,
withloop=withloop,
activation=activation,
dropout=dropout,
dense=True,
aggrmethod=aggrmethod,
)
def forward(self, graph, x):
return self.model.forward(graph, x)
def get_outdim(self):
return self.model.get_outdim()
def __repr__(self):
return "%s %s (%d - [%d:%d] > %d)" % (
self.__class__.__name__,
self.aggrmethod,
self.model.in_features,
self.model.hiddendim,
self.model.nhiddenlayer,
self.model.out_features,
)
class InceptionGCNBlock(Module):
"""
The multiple layer GCN with inception connection block.
"""
def __init__(
self,
in_features,
out_features,
nbaselayer,
withbn=True,
withloop=True,
activation=F.relu,
dropout=True,
aggrmethod="concat",
dense=False,
):
"""
The multiple layer GCN with inception connection block.
:param in_features: the input feature dimension.
:param out_features: the hidden feature dimension.
:param nbaselayer: the number of layers in the base block.
:param withbn: using batch normalization in graph convolution.
:param withloop: using self feature modeling in graph convolution.
:param activation: the activation function, default is ReLu.
:param dropout: the dropout ratio.
:param aggrmethod: the aggregation function for baseblock, can be "concat" and "add". For "resgcn", the default
is "add", for others the default is "concat".
:param dense: not applied. The default is False, cannot be changed.
"""
super(InceptionGCNBlock, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.hiddendim = out_features
self.nbaselayer = nbaselayer
self.activation = activation
self.aggrmethod = aggrmethod
self.dropout = dropout
self.withbn = withbn
self.withloop = withloop
self.midlayers = nn.ModuleList()
self.__makehidden()
if self.aggrmethod == "concat":
self.out_features = in_features + out_features * nbaselayer
elif self.aggrmethod == "add":
if in_features != self.hiddendim:
raise RuntimeError("The dimension of in_features and hiddendim should be matched in 'add' model.")
self.out_features = out_features
else:
raise NotImplementedError("The aggregation method only support 'concat', 'add'.")
def __makehidden(self):
# for j in xrange(self.nhiddenlayer):
for j in range(self.nbaselayer):
reslayer = nn.ModuleList()
# for i in xrange(j + 1):
for i in range(j + 1):
if i == 0:
layer = GraphConvolutionBS(
self.in_features,
self.hiddendim,
self.activation,
self.withbn,
self.withloop,
)
else:
layer = GraphConvolutionBS(
self.hiddendim,
self.hiddendim,
self.activation,
self.withbn,
self.withloop,
)
reslayer.append(layer)
self.midlayers.append(reslayer)
def forward(self, graph, x):
for reslayer in self.midlayers:
subx = x
for gc in reslayer:
subx = gc(graph, x)
subx = F.dropout(subx, self.dropout, training=self.training)
x = self._doconcat(x, subx)
return x
def get_outdim(self):
return self.out_features
def _doconcat(self, x, subx):
if self.aggrmethod == "concat":
return torch.cat((x, subx), 1)
elif self.aggrmethod == "add":
return x + subx
def __repr__(self):
return "%s %s (%d - [%d:%d] > %d)" % (
self.__class__.__name__,
self.aggrmethod,
self.in_features,
self.hiddendim,
self.nbaselayer,
self.out_features,
)
class Dense(Module):
"""
Simple Dense layer, Do not consider adj.
"""
def __init__(self, in_features, out_features, activation=lambda x: x, bias=True, res=False):
super(Dense, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.sigma = activation
self.weight = Parameter(torch.FloatTensor(in_features, out_features))
self.res = res
self.bn = nn.BatchNorm1d(out_features)
self.bias = Parameter(torch.FloatTensor(out_features)) if bias else None
self.reset_parameters()
def reset_parameters(self):
stdv = 1.0 / math.sqrt(self.weight.size(1))
self.weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform_(-stdv, stdv)
def forward(self, graph, x):
output = torch.mm(x, self.weight)
output = output + self.bias if self.bias is not None else output
output = self.bn(output)
return self.sigma(output)
def __repr__(self):
return self.__class__.__name__ + " (" + str(self.in_features) + " -> " + str(self.out_features) + ")"
[docs]@register_model("dropedge_gcn")
class DropEdge_GCN(BaseModel):
"""
DropEdge: Towards Deep Graph Convolutional Networks on Node Classification
Applying DropEdge to GCN @ https://arxiv.org/pdf/1907.10903.pdf
The model for the single kind of deepgcn blocks.
The model architecture likes:
inputlayer(nfeat)--block(nbaselayer, nhid)--...--outputlayer(nclass)--softmax(nclass)
|------ nhidlayer ----|
The total layer is nhidlayer*nbaselayer + 2.
All options are configurable.
Args:
Initial function.
:param nfeat: the input feature dimension.
:param nhid: the hidden feature dimension.
:param nclass: the output feature dimension.
:param nhidlayer: the number of hidden blocks.
:param dropout: the dropout ratio.
:param baseblock: the baseblock type, can be "mutigcn", "resgcn", "densegcn" and "inceptiongcn".
:param inputlayer: the input layer type, can be "gcn", "dense", "none".
:param outputlayer: the input layer type, can be "gcn", "dense".
:param nbaselayer: the number of layers in one hidden block.
:param activation: the activation function, default is ReLu.
:param withbn: using batch normalization in graph convolution.
:param withloop: using self feature modeling in graph convolution.
:param aggrmethod: the aggregation function for baseblock, can be "concat" and "add". For "resgcn", the default
is "add", for others the default is "concat".
"""
[docs] @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('--baseblock', default='mutigcn',
help="Choose the model to be trained.( mutigcn, resgcn, densegcn, inceptiongcn)")
parser.add_argument('--inputlayer', default='gcn',
help="The input layer of the model.")
parser.add_argument('--outputlayer', default='gcn',
help="The output layer of the model.")
parser.add_argument('--hidden_size', type=int, default=64,
help='Number of hidden units.')
parser.add_argument('--dropout', type=float, default=0.5,
help='Dropout rate (1 - keep probability).')
parser.add_argument('--withbn', action='store_true', default=False,
help='Enable Bath Norm GCN')
parser.add_argument('--withloop', action="store_true", default=False,
help="Enable loop layer GCN")
parser.add_argument('--nhiddenlayer', type=int, default=1,
help='The number of hidden layers.')
parser.add_argument("--nbaseblocklayer", type=int, default=0,
help="The number of layers in each baseblock")
parser.add_argument("--aggrmethod", default="default",
help="The aggrmethod for the layer aggreation. The options includes add and concat. Only valid in resgcn, densegcn and inceptiongcn")
parser.add_argument("--task_type", default="full", help="The node classification task type (full and semi). Only valid for cora, citeseer and pubmed dataset.")
parser.add_argument("--activation", default=F.relu, help="activiation function")
# fmt: on
[docs] @classmethod
def build_model_from_args(cls, args):
return cls(
args.num_features,
args.hidden_size,
args.num_classes,
args.nhiddenlayer,
args.dropout,
args.baseblock,
args.inputlayer,
args.outputlayer,
args.nbaseblocklayer,
args.activation,
args.withbn,
args.withloop,
args.aggrmethod,
)
def __init__(
self,
nfeat,
nhid,
nclass,
nhidlayer,
dropout,
baseblock,
inputlayer,
outputlayer,
nbaselayer,
activation,
withbn,
withloop,
aggrmethod,
):
super(DropEdge_GCN, self).__init__()
self.dropout = dropout
if baseblock == "resgcn":
self.BASEBLOCK = ResGCNBlock
elif baseblock == "densegcn":
self.BASEBLOCK = DenseGCNBlock
elif baseblock == "mutigcn":
self.BASEBLOCK = MultiLayerGCNBlock
elif baseblock == "inceptiongcn":
self.BASEBLOCK = InceptionGCNBlock
else:
raise NotImplementedError("Current baseblock %s is not supported." % (baseblock))
if inputlayer == "gcn":
# input gc
self.ingc = GraphConvolutionBS(nfeat, nhid, activation, withbn, withloop)
baseblockinput = nhid
elif inputlayer == "none":
self.ingc = lambda x: x
baseblockinput = nfeat
else:
self.ingc = Dense(nfeat, nhid, activation)
baseblockinput = nhid
outactivation = lambda x: x # noqa E731
if outputlayer == "gcn":
self.outgc = GraphConvolutionBS(baseblockinput, nclass, outactivation, withbn, withloop)
# elif outputlayer == "none": #here can not be none
# self.outgc = lambda x: x
else:
self.outgc = Dense(nhid, nclass, activation)
# hidden layer
self.midlayer = nn.ModuleList()
# Dense is not supported now.
# for i in xrange(nhidlayer):
for i in range(nhidlayer):
gcb = self.BASEBLOCK(
in_features=baseblockinput,
out_features=nhid,
nbaselayer=nbaselayer,
withbn=withbn,
withloop=withloop,
activation=activation,
dropout=dropout,
dense=False,
aggrmethod=aggrmethod,
)
self.midlayer.append(gcb)
baseblockinput = gcb.get_outdim()
# output gc
outactivation = lambda x: x # noqa E731 we donot need nonlinear activation here.
self.outgc = GraphConvolutionBS(baseblockinput, nclass, outactivation, withbn, withloop)
self.reset_parameters()
[docs] def reset_parameters(self):
pass
[docs] def forward(self, graph):
x = graph.x
x = self.ingc(graph, x)
x = F.dropout(x, self.dropout, training=self.training)
# mid block connections
# for i in xrange(len(self.midlayer)):
for i in range(len(self.midlayer)):
midgc = self.midlayer[i]
x = midgc(graph, x)
# output, no relu and dropput here.
x = self.outgc(graph, x)
x = F.log_softmax(x, dim=1)
return x
[docs] def predict(self, data):
return self.forward(data)