def __init__(self,
in_channels,
hidden_channels,
out_channels,
num_layers,
dropout,
mode='concat'):
super().__init__()
self.convs = torch.nn.ModuleList()
self.convs.append(GCNConv(in_channels, hidden_channels, cached=False))
self.bns = torch.nn.ModuleList()
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
for _ in range(num_layers - 1):
self.convs.append(
GCNConv(hidden_channels, hidden_channels, cached=False))
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = Linear(num_layers * hidden_channels, hidden_channels)
else:
self.lin1 = Linear(hidden_channels, hidden_channels)
self.lin2 = Linear(hidden_channels, out_channels)
self.dropout = dropout
def __init__(self,
nfeat,
nhid,
nclass,
dropout=0.5,
lr=0.01,
weight_decay=5e-4,
n_edge=1,
with_relu=True,
drop=False,
with_bias=True,
device=None):
super(GIN, self).__init__()
assert device is not None, "Please specify 'device'!"
self.device = device
self.nfeat = nfeat
self.hidden_sizes = [nhid]
self.nclass = int(nclass)
self.dropout = dropout
self.lr = lr
self.drop = drop
if not with_relu:
self.weight_decay = 0
else:
self.weight_decay = weight_decay
self.with_relu = with_relu
self.with_bias = with_bias
self.n_edge = n_edge
self.output = None
self.best_model = None
self.best_output = None
self.adj_norm = None
self.features = None
self.gate = Parameter(torch.rand(1)) # creat a generator between [0,1]
nclass = int(nclass)
"""GIN from torch-geometric"""
num_features = nfeat
dim = nhid
nn1 = Sequential(
Linear(num_features, dim),
ReLU(),
)
self.gc1 = GINConv(nn1)
# self.bn1 = torch.nn.BatchNorm1d(dim)
nn2 = Sequential(
Linear(dim, dim),
ReLU(),
)
self.gc2 = GINConv(nn2)
nn3 = Sequential(
Linear(dim, dim),
ReLU(),
)
self.gc3 = GINConv(nn3)
self.jump = JumpingKnowledge(mode='cat')
# self.bn2 = torch.nn.BatchNorm1d(dim)
self.fc1 = Linear(dim, dim)
self.fc2 = Linear(dim * 1, int(nclass))
def __init__(self, dataset, num_layers, hidden, train_eps=False, mode='cat'):
super().__init__()
self.conv1 = GINConv(nn.Sequential(
nn.Linear(dataset.num_features, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.BatchNorm1d(hidden),
), train_eps=train_eps)
self.convs = nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(
GINConv(nn.Sequential(
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.BatchNorm1d(hidden),
), train_eps=train_eps))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = nn.Linear(num_layers * hidden, hidden)
else:
self.lin1 = nn.Linear(hidden, hidden)
self.lin2 = nn.Linear(hidden, dataset.num_classes)
def __init__(self,
args,
num_nodes=10,
num_layers=4,
hidden=16,
ratio=0.25):
super(DiffPool, self).__init__()
self.args = args
num_features = self.args.filters_3
self.att = DenseAttentionModule(self.args)
num_nodes = ceil(ratio * num_nodes)
self.embed_block1 = Block(num_features, hidden, hidden)
self.pool_block1 = Block(num_features, hidden, num_nodes)
self.embed_blocks = torch.nn.ModuleList()
self.pool_blocks = torch.nn.ModuleList()
for i in range((num_layers // 2) - 1):
num_nodes = ceil(ratio * num_nodes)
self.embed_blocks.append(Block(hidden, hidden, hidden))
self.pool_blocks.append(Block(hidden, hidden, num_nodes))
self.jump = JumpingKnowledge(mode="cat")
self.lin1 = Linear((len(self.embed_blocks) + 1) * hidden, hidden)
self.lin2 = Linear(hidden, num_features)
class DiffPool(torch.nn.Module):
def __init__(self,
args,
num_nodes=10,
num_layers=4,
hidden=16,
ratio=0.25):
super(DiffPool, self).__init__()
self.args = args
num_features = self.args.filters_3
self.att = DenseAttentionModule(self.args)
num_nodes = ceil(ratio * num_nodes)
self.embed_block1 = Block(num_features, hidden, hidden)
self.pool_block1 = Block(num_features, hidden, num_nodes)
self.embed_blocks = torch.nn.ModuleList()
self.pool_blocks = torch.nn.ModuleList()
for i in range((num_layers // 2) - 1):
num_nodes = ceil(ratio * num_nodes)
self.embed_blocks.append(Block(hidden, hidden, hidden))
self.pool_blocks.append(Block(hidden, hidden, num_nodes))
self.jump = JumpingKnowledge(mode="cat")
self.lin1 = Linear((len(self.embed_blocks) + 1) * hidden, hidden)
self.lin2 = Linear(hidden, num_features)
def reset_parameters(self):
self.embed_block1.reset_parameters()
self.pool_block1.reset_parameters()
for block1, block2 in zip(self.embed_blocks, self.pool_blocks):
block1.reset_parameters()
block2.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, x, adj, mask):
s = self.pool_block1(x, adj, mask, add_loop=True)
x = F.relu(self.embed_block1(x, adj, mask, add_loop=True))
xs = [self.att(x, mask)]
x, adj, _, _ = dense_diff_pool(x, adj, s, mask)
for i, (embed,
pool) in enumerate(zip(self.embed_blocks, self.pool_blocks)):
s = pool(x, adj)
x = F.relu(embed(x, adj))
xs.append(self.att(x))
if i < (len(self.embed_blocks) - 1):
x, adj, _, _ = dense_diff_pool(x, adj, s)
x = self.jump(xs)
x = F.relu(self.lin1(x))
x = self.lin2(x)
return x
def __repr__(self):
return self.__class__.__name__
def __init__(self, hparams: dict):
super().__init__()
self.hparams = hparams
self.aggregation_method = hparams["aggregation_method"]
if hparams["num_conv_layers"] < 1:
raise Exception("Invalid number of layers!")
self.conv_modules = nn.ModuleList()
self.conv_modules.append(
GCNConv(hparams["num_node_features"], hparams["conv_size"]))
for _ in range(hparams["num_conv_layers"] - 1):
conv = GCNConv(hparams["conv_size"], hparams["conv_size"])
self.conv_modules.append(conv)
if self.aggregation_method == 'lstm':
self.jk = JumpingKnowledge(self.aggregation_method,
num_layers=hparams["num_conv_layers"],
channels=hparams["conv_size"])
else:
self.jk = JumpingKnowledge(self.aggregation_method)
if self.aggregation_method == 'cat':
self.lin = nn.Linear(
int(hparams["conv_size"] * hparams["num_conv_layers"]),
hparams["lin_size"])
else:
self.lin = nn.Linear(int(hparams["conv_size"]),
hparams["lin_size"])
self.output = nn.Linear(hparams["lin_size"], hparams["output_size"])
class DiffPool(torch.nn.Module):
def __init__(self, dataset, num_pools, hidden, ratio=0.25):
super(DiffPool, self).__init__()
self.num_pools, self.hidden = num_pools, hidden
num_nodes = ceil(ratio * dataset[0].num_nodes)
self.embed_block1 = Block(dataset.num_features, hidden, hidden)
self.pool_block1 = Block(dataset.num_features, hidden, num_nodes)
self.embed_blocks = torch.nn.ModuleList()
self.pool_blocks = torch.nn.ModuleList()
for i in range(num_pools - 1):
num_nodes = ceil(ratio * num_nodes)
self.embed_blocks.append(Block(hidden, hidden, hidden))
self.pool_blocks.append(Block(hidden, hidden, num_nodes))
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear((len(self.embed_blocks) + 1) * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def reset_parameters(self):
self.embed_block1.reset_parameters()
self.pool_block1.reset_parameters()
for embed_block, pool_block in zip(self.embed_blocks,
self.pool_blocks):
embed_block.reset_parameters()
pool_block.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
# x:[batch_size,num_nodes,in_channels]
x, adj, mask = data.x, data.adj, data.mask
# x:[batch_size, num_nodes, c_num_nodes]
s = self.pool_block1(x, adj, mask, add_loop=True)
# s:[batch_size, num_nodes, hidden]
x = F.relu(self.embed_block1(x, adj, mask, add_loop=True))
xs = [x.mean(dim=1)]
# x:[batch_size, c_num_nodes, hidden]
x, adj, _, _ = dense_diff_pool(x, adj, s, mask)
# adj: [batch_size,c_num_nodes, c_num_nodes]
for i, (embed_block, pool_block) in enumerate(
zip(self.embed_blocks, self.pool_blocks)):
# s: [batch_size,c_num_nodes, cc_num_nodes]
s = pool_block(x, adj)
# x: [batch_size,c_num_nodes,hidden]
x = F.relu(embed_block(x, adj))
xs.append(x.mean(dim=1))
if i < len(self.embed_blocks) - 1:
# x: [batch_size,cc_num_nodes, hidden]
x, adj, _, _ = dense_diff_pool(x, adj, s)
# adj: [batch_size,cc_num_nodes,cc_num_nodes]
x = self.jump(xs) # x: [batch_size,len(self.embed_blocks)+1)*hidden]
x = F.relu(self.lin1(x)) # x: [batch_size,hidden]
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x) # x: [batch_size,dataset.num_classes]
return F.log_softmax(x, dim=-1)
def __init__(self, dataset, num_layers, hidden, mode='cat'):
super(GIN0WithJK, self).__init__()
self.conv1 = GINConv(Sequential(
Linear(dataset.num_features, hidden),
ReLU(),
Linear(hidden, hidden),
ReLU(),
BN(hidden),
),
train_eps=False)
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(
GINConv(Sequential(
Linear(hidden, hidden),
ReLU(),
Linear(hidden, hidden),
ReLU(),
BN(hidden),
),
train_eps=False))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = Linear(num_layers * hidden, hidden)
else:
self.lin1 = Linear(hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2,
dropout=0.5, heads=2, jk_type='max'):
super(GATJK, self).__init__()
self.convs = nn.ModuleList()
self.convs.append(
GATConv(in_channels, hidden_channels, heads=heads, concat=True))
self.bns = nn.ModuleList()
self.bns.append(nn.BatchNorm1d(hidden_channels*heads))
for _ in range(num_layers - 2):
self.convs.append(
GATConv(hidden_channels*heads, hidden_channels, heads=heads, concat=True) )
self.bns.append(nn.BatchNorm1d(hidden_channels*heads))
self.convs.append(
GATConv(hidden_channels*heads, hidden_channels, heads=heads))
self.dropout = dropout
self.activation = F.elu # note: uses elu
self.jump = JumpingKnowledge(jk_type, channels=hidden_channels*heads, num_layers=1)
if jk_type == 'cat':
self.final_project = nn.Linear(hidden_channels*heads*num_layers, out_channels)
else: # max or lstm
self.final_project = nn.Linear(hidden_channels*heads, out_channels)
def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2,
dropout=0.5, save_mem=False, jk_type='max'):
super(GCNJK, self).__init__()
self.convs = nn.ModuleList()
self.convs.append(
GCNConv(in_channels, hidden_channels, cached=not save_mem, normalize=not save_mem))
self.bns = nn.ModuleList()
self.bns.append(nn.BatchNorm1d(hidden_channels))
for _ in range(num_layers - 2):
self.convs.append(
GCNConv(hidden_channels, hidden_channels, cached=not save_mem, normalize=not save_mem))
self.bns.append(nn.BatchNorm1d(hidden_channels))
self.convs.append(
GCNConv(hidden_channels, hidden_channels, cached=not save_mem, normalize=not save_mem))
self.dropout = dropout
self.activation = F.relu
self.jump = JumpingKnowledge(jk_type, channels=hidden_channels, num_layers=1)
if jk_type == 'cat':
self.final_project = nn.Linear(hidden_channels * num_layers, out_channels)
else: # max or lstm
self.final_project = nn.Linear(hidden_channels, out_channels)
class GCNWithJK(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden):
super(GCNWithJK, self).__init__()
self.conv1 = GCNConv(dataset.num_features, hidden)
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(GCNConv(hidden, hidden))
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(num_layers * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def reset_parameters(self):
self.conv1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = F.relu(self.conv1(x, edge_index))
xs = [x]
for conv in self.convs:
x = F.relu(conv(x, edge_index))
xs += [x]
x = self.jump(xs)
x = global_mean_pool(x, batch)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
return F.log_softmax(x, dim=-1)
def __repr__(self):
return self.__class__.__name__
class Coarsening(torch.nn.Module):
def __init__(self, dataset, hidden, ratio=0.25): # we only use 1 layer for coarsening
super(Coarsening, self).__init__()
# self.embed_block1 = GNNBlock(dataset.num_features, hidden, hidden)
self.embed_block1 = DenseGCNConv(dataset.num_features, hidden)
self.coarse_block1 = CoarsenBlock(hidden, ratio)
self.embed_block2 = DenseGCNConv(hidden, dataset.num_features)
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(hidden + dataset.num_features, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def reset_parameters(self):
self.embed_block1.reset_parameters()
self.coarse_block1.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data, epsilon=0.01, opt_epochs=100):
x, adj, mask = data.x, data.adj, data.mask
batch_num_nodes = data.mask.sum(-1)
x1 = F.relu(self.embed_block1(x, adj, mask, add_loop=True))
# xs = [x1.mean(dim=1)]
coarse_x, new_adj, S = self.coarse_block1(x1, adj, batch_num_nodes)
xs = [coarse_x.mean(dim=1)]
x2 = F.tanh(self.embed_block2(coarse_x, new_adj, mask, add_loop=True))
xs.append(x2.mean(dim=1))
opt_loss = 0.0
for i in range(len(x)):
x3 = self.get_nonzero_rows(x[i])
x4 = self.get_nonzero_rows(x2[i])
# if x3.size()[0]==0 or x4.size()[0]==0:
# continue
# opt_loss += sinkhorn_loss_default(x3, x4, epsilon, niter=opt_epochs).float()
opt_loss += sinkhorn_loss_default(x3, x2[i], epsilon, niter=opt_epochs)
return xs, new_adj, S, opt_loss
def predict(self, xs):
x = self.jump(xs)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
return F.log_softmax(x, dim=-1)
def get_nonzero_rows(self, M):# M is a matrix
# row_ind = M.sum(-1).nonzero().squeeze() #nonzero has bugs in Pytorch 1.2.0.........
#So we use other methods to take place of it
MM, MM_ind = M.sum(-1).sort()
N = (M.sum(-1)>0).sum()
return M[MM_ind[:N]]
def __repr__(self):
return self.__class__.__name__
class DiffPool(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden, ratio=0.25):
super(DiffPool, self).__init__()
num_nodes = ceil(ratio * dataset[0].num_nodes)
self.embed_block1 = Block(dataset.num_features, hidden, hidden)
self.pool_block1 = Block(dataset.num_features, hidden, num_nodes)
self.embed_blocks = torch.nn.ModuleList()
self.pool_blocks = torch.nn.ModuleList()
for i in range((num_layers // 2) - 1):
num_nodes = ceil(ratio * num_nodes)
self.embed_blocks.append(Block(hidden, hidden, hidden))
self.pool_blocks.append(Block(hidden, hidden, num_nodes))
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear((len(self.embed_blocks) + 1) * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def reset_parameters(self):
self.embed_block1.reset_parameters()
self.pool_block1.reset_parameters()
for embed_block, pool_block in zip(self.embed_blocks,
self.pool_blocks):
embed_block.reset_parameters()
pool_block.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
x, adj, mask = data.x, data.adj, data.mask
link_losses = 0.
ent_losses = 0.
s = self.pool_block1(x, adj, mask, add_loop=True)
x = F.relu(self.embed_block1(x, adj, mask, add_loop=True))
xs = [x.mean(dim=1)]
x, adj, link_loss, ent_loss = dense_diff_pool(x, adj, s, mask)
link_losses += link_loss
ent_losses += ent_loss
for i, (embed_block, pool_block) in enumerate(
zip(self.embed_blocks, self.pool_blocks)):
s = pool_block(x, adj)
x = F.relu(embed_block(x, adj))
xs.append(x.mean(dim=1))
if i < len(self.embed_blocks) - 1:
x, adj, link_loss, ent_loss = dense_diff_pool(x, adj, s)
link_losses += link_loss
ent_losses += ent_loss
x = self.jump(xs)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
return F.log_softmax(x, dim=-1), link_losses + ent_losses
def __repr__(self):
return self.__class__.__name__
def __init__(self, dataset, num_layers, hidden):
super(Graclus, self).__init__()
self.conv1 = GraphConv(dataset.num_features, hidden, aggr='mean')
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(GraphConv(hidden, hidden, aggr='mean'))
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(num_layers * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def test_jumping_knowledge():
num_nodes, channels, num_layers = 100, 17, 5
xs = list([torch.randn(num_nodes, channels) for _ in range(num_layers)])
model = JumpingKnowledge('cat')
assert model.__repr__() == 'JumpingKnowledge(cat)'
out = model(xs)
assert out.size() == (num_nodes, channels * num_layers)
if is_full_test():
jit = torch.jit.script(model)
assert torch.allclose(jit(xs), out)
model = JumpingKnowledge('max')
assert model.__repr__() == 'JumpingKnowledge(max)'
out = model(xs)
assert out.size() == (num_nodes, channels)
if is_full_test():
jit = torch.jit.script(model)
assert torch.allclose(jit(xs), out)
model = JumpingKnowledge('lstm', channels, num_layers)
assert model.__repr__() == 'JumpingKnowledge(lstm)'
out = model(xs)
assert out.size() == (num_nodes, channels)
if is_full_test():
jit = torch.jit.script(model)
assert torch.allclose(jit(xs), out)
def __init__(self, num_input_features, num_layers, hidden, mode='cat'):
super(GraphSAGEWithJK, self).__init__()
self.conv1 = SAGEConv(num_input_features, hidden)
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(SAGEConv(hidden, hidden))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = Linear(3 * num_layers * hidden, hidden)
else:
self.lin1 = Linear(3 * hidden, hidden)
self.lin2 = Linear(hidden, 2)
class GATJK(nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2,
dropout=0.5, heads=2, jk_type='max'):
super(GATJK, self).__init__()
self.convs = nn.ModuleList()
self.convs.append(
GATConv(in_channels, hidden_channels, heads=heads, concat=True))
self.bns = nn.ModuleList()
self.bns.append(nn.BatchNorm1d(hidden_channels*heads))
for _ in range(num_layers - 2):
self.convs.append(
GATConv(hidden_channels*heads, hidden_channels, heads=heads, concat=True) )
self.bns.append(nn.BatchNorm1d(hidden_channels*heads))
self.convs.append(
GATConv(hidden_channels*heads, hidden_channels, heads=heads))
self.dropout = dropout
self.activation = F.elu # note: uses elu
self.jump = JumpingKnowledge(jk_type, channels=hidden_channels*heads, num_layers=1)
if jk_type == 'cat':
self.final_project = nn.Linear(hidden_channels*heads*num_layers, out_channels)
else: # max or lstm
self.final_project = nn.Linear(hidden_channels*heads, out_channels)
def reset_parameters(self):
for conv in self.convs:
conv.reset_parameters()
for bn in self.bns:
bn.reset_parameters()
self.jump.reset_parameters()
self.final_project.reset_parameters()
def forward(self, data):
x = data.graph['node_feat']
xs = []
for i, conv in enumerate(self.convs[:-1]):
x = conv(x, data.graph['edge_index'])
x = self.bns[i](x)
x = self.activation(x)
xs.append(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.convs[-1](x, data.graph['edge_index'])
xs.append(x)
x = self.jump(xs)
x = self.final_project(x)
return x
def __init__(self, dataset, hidden, ratio=0.25): # we only use 1 layer for coarsening
super(Coarsening, self).__init__()
# self.embed_block1 = GNNBlock(dataset.num_features, hidden, hidden)
self.embed_block1 = DenseGCNConv(dataset.num_features, hidden)
self.coarse_block1 = CoarsenBlock(hidden, ratio)
self.embed_block2 = DenseGCNConv(hidden, dataset.num_features)
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(hidden + dataset.num_features, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def __init__(self, dataset, num_layers, hidden, mode='cat'):
super(GCNWithJK, self).__init__()
self.conv1 = GCNConv(dataset.num_features, hidden)
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(GCNConv(hidden, hidden))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = Linear(num_layers * hidden, hidden)
else:
self.lin1 = Linear(hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def __init__(self, dataset, hidden, num_layers=2, ratio=0.5):
super(MultiLayerCoarsening, self).__init__()
self.embed_block1 = DenseGCNConv(dataset.num_features, hidden)
self.coarse_block1 = CoarsenBlock(hidden, ratio)
self.embed_block2 = DenseGCNConv(hidden, dataset.num_features)
# self.embed_block2 = GNNBlock(hidden, hidden, dataset.num_features)
self.num_layers = num_layers
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(hidden + dataset.num_features, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
class JKNet(torch.nn.Module):
def __init__(self,
in_channels,
hidden_channels,
out_channels,
num_layers,
dropout,
mode='concat'):
super().__init__()
self.convs = torch.nn.ModuleList()
self.convs.append(GCNConv(in_channels, hidden_channels, cached=False))
self.bns = torch.nn.ModuleList()
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
for _ in range(num_layers - 1):
self.convs.append(
GCNConv(hidden_channels, hidden_channels, cached=False))
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = Linear(num_layers * hidden_channels, hidden_channels)
else:
self.lin1 = Linear(hidden_channels, hidden_channels)
self.lin2 = Linear(hidden_channels, out_channels)
self.dropout = dropout
def reset_parameters(self):
for conv in self.convs:
conv.reset_parameters()
for bn in self.bns:
bn.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, x, adj_t):
xs = []
for i, conv in enumerate(self.convs):
x = conv(x, adj_t)
x = self.bns[i](x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
xs += [x]
x = self.jump(xs)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.lin2(x)
return F.log_softmax(x, dim=-1)
class DiffPool(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden):
super(DiffPool, self).__init__()
num_nodes = ceil(0.25 * dataset[0].num_nodes)
self.embed_block1 = Block(dataset.num_features, hidden, hidden)
self.pool_block1 = Block(dataset.num_features, hidden, num_nodes)
self.embed_blocks = torch.nn.ModuleList()
self.pool_blocks = torch.nn.ModuleList()
for i in range((num_layers // 2) - 1):
num_nodes = ceil(0.25 * num_nodes)
self.embed_blocks.append(Block(hidden, hidden, hidden))
self.pool_blocks.append(Block(hidden, hidden, num_nodes))
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear((len(self.embed_blocks) + 1) * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
def reset_parameters(self):
self.embed_block1.reset_parameters()
self.pool_block1.reset_parameters()
for block1, block2 in zip(self.embed_blocks, self.pool_blocks):
block1.reset_parameters()
block2.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
x, adj, mask = data.x, data.adj, data.mask
s = self.pool_block1(x, adj, mask, add_loop=True)
x = F.relu(self.embed_block1(x, adj, mask, add_loop=True))
xs = [x.mean(dim=1)]
x, adj, reg = dense_diff_pool(x, adj, s, mask)
for embed, pool in zip(self.embed_blocks, self.pool_blocks):
s = pool(x, adj)
x = F.relu(embed(x, adj))
xs.append(x.mean(dim=1))
x, adj, _, _ = dense_diff_pool(x, adj, s)
x = self.jump(xs)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
return F.log_softmax(x, dim=-1)
def __repr__(self):
return self.__class__.__name__
class GIN(nn.Module):
def __init__(self, dataset, num_layers, hidden, train_eps=False, mode='cat'):
super().__init__()
self.conv1 = GINConv(nn.Sequential(
nn.Linear(dataset.num_features, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.BatchNorm1d(hidden),
), train_eps=train_eps)
self.convs = nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(
GINConv(nn.Sequential(
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.BatchNorm1d(hidden),
), train_eps=train_eps))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = nn.Linear(num_layers * hidden, hidden)
else:
self.lin1 = nn.Linear(hidden, hidden)
self.lin2 = nn.Linear(hidden, dataset.num_classes)
def reset_parameters(self):
self.conv1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
self.jump.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x = self.conv1(x, edge_index)
xs = [x]
for conv in self.convs:
x = conv(x, edge_index)
xs += [x]
x = self.jump(xs)
x = global_mean_pool(x, batch)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
return x
def __repr__(self):
return self.__class__.__name__
class ModelEnsemble(torch.nn.Module):
def __init__(self, num_features, num_class):
super(ModelEnsemble, self).__init__()
self.JK = JumpingKnowledge(mode='lstm',
channels=num_class,
num_layers=1)
# self.linear = Linear(num_features, num_class)
def reset_parameters(self):
self.JK.reset_parameters()
def forward(self, x):
x = self.JK(x)
return log_softmax(x, dim=-1)
def __init__(self, num_vocab, max_seq_len, node_encoder, emb_dim, num_layers, hidden, ratio=0.8, dropout=0, num_class=0):
super(ASAP, self).__init__()
self.num_class = num_class
self.max_seq_len = max_seq_len
self.node_encoder = node_encoder
self.conv1 = GraphConv(emb_dim, hidden, aggr='mean')
self.convs = torch.nn.ModuleList()
self.pools = torch.nn.ModuleList()
self.convs.extend([
GraphConv(hidden, hidden, aggr='mean')
for i in range(num_layers - 1)
])
self.pools.extend([
ASAPooling(hidden, ratio, dropout=dropout)
for i in range((num_layers) // 2)
])
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(num_layers * hidden, hidden)
# self.lin2 = Linear(hidden, dataset.num_classes)
if self.num_class > 0: # classification
self.graph_pred_linear = torch.nn.Linear(hidden, self.num_class)
else:
self.graph_pred_linear_list = torch.nn.ModuleList()
for i in range(max_seq_len):
self.graph_pred_linear_list.append(torch.nn.Linear(hidden, num_vocab))
def __init__(self, num_layers, hidden_list, activation, data):
super(ModelGCN, self).__init__()
assert len(hidden_list) == num_layers + 1
self.linear_1 = Linear(data.num_features, hidden_list[0])
self.convs = torch.nn.ModuleList()
for i in range(num_layers):
self.convs.append(GCNConv(hidden_list[i], hidden_list[i + 1]))
self.JK = JumpingKnowledge(mode='max')
self.linear_2 = Linear(hidden_list[-1], data.num_class)
if activation == "relu":
self.activation = relu
elif activation == "leaky_relu":
self.activation = leaky_relu
self.reg_params = list(self.linear_1.parameters()) + list(
self.convs.parameters()) + list(self.JK.parameters()) + list(
self.linear_2.parameters())
def __init__(self, dataset, num_layers, hidden, jpgs=False, jp=False, hop=2, num_patches=5, \
ratio=0.25, plot=False, dropout=False, ge=False):
super(DiffPool, self).__init__()
Block = [Block_1hop, Block_2hop, Block_3hop][hop-1]
self.num_patches = num_patches
self.dropout = dropout
self.plot = plot
num_nodes = ceil(ratio * dataset[0].num_nodes)
self.embed_block1 = Block(dataset.num_features, hidden, hidden, jpgs)
self.pool_block1 = Block(dataset.num_features, hidden, num_nodes, jpgs)
self.embed_blocks = torch.nn.ModuleList()
self.pool_blocks = torch.nn.ModuleList()
for i in range((num_layers // 2) - 1):
num_nodes = ceil(ratio * num_nodes)
self.embed_blocks.append(Block(hidden, hidden, hidden, jpgs))
self.pool_blocks.append(Block(hidden, hidden, num_nodes, jpgs))
self.pool_block_last = Block(hidden, hidden, 1, jpgs)
self.jp = jp
self.ge = ge
if self.jp:
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear((len(self.embed_blocks) + 1) * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
else:
# self.lin1 = Linear(num_patches*hidden, hidden)
# !!! 因为改成了average graph embeddings，所以dimension也需要修改
self.lin1 = Linear(hidden, dataset.num_classes)
def __init__(self,
in_channels,
out_channels,
depth=3,
jk_depth=7,
base_conv=TAGConv,
att_conv=GATConv,
base_conv_settings={"K": 3},
att_conv_settings={}):
super(AttentionBlock, self).__init__()
assert out_channels % in_channels == 0, "out_channels must be a multiple of in_channels"
ratio = out_channels // in_channels
att_conv_settings = dict(**att_conv_settings, heads=ratio)
self.base_conv_list = torch.nn.ModuleList([])
self.att_conv_list = torch.nn.ModuleList([])
self.slc_list = torch.nn.ModuleList([])
self.jk = JumpingKnowledge("lstm", ratio * in_channels, jk_depth)
for _ in range(depth):
self.slc_list.append(SublayerConnection(in_channels))
self.base_conv_list.append(
base_conv(in_channels, in_channels, **base_conv_settings))
self.att_conv_list.append(
att_conv(in_channels, in_channels, **att_conv_settings))
self.reset_parameters()
def __init__(self, in_channels, hidden_channels, out_channels, mode="cat"):
super(Block, self).__init__()
# self.conv1 = DenseSAGEConv(in_channels, hidden_channels)
# self.conv2 = DenseSAGEConv(hidden_channels, out_channels)
# self.conv1 = DenseGCNConv(in_channels, hidden_channels)
# self.conv2 = DenseGCNConv(hidden_channels, out_channels)
nn1 = torch.nn.Sequential(
Linear(in_channels, hidden_channels),
ReLU(),
Linear(hidden_channels, hidden_channels),
)
nn2 = torch.nn.Sequential(
Linear(hidden_channels, out_channels),
ReLU(),
Linear(out_channels, out_channels),
)
self.conv1 = DenseGINConv(nn1, train_eps=True)
self.conv2 = DenseGINConv(nn2, train_eps=True)
self.jump = JumpingKnowledge(mode)
if mode == "cat":
self.lin = Linear(hidden_channels + out_channels, out_channels)
else:
self.lin = Linear(out_channels, out_channels)
def __init__(self,
dataset,
num_layers,
hidden,
weight_conv='WeightConv1',
multi_channel='False'):
super(SMG_JK, self).__init__()
self.lin0 = Linear(dataset.num_features, hidden)
self.convs = torch.nn.ModuleList()
for i in range(num_layers):
self.convs.append(SparseConv(hidden, hidden))
self.masks = torch.nn.ModuleList()
if multi_channel == 'True':
out_channel = hidden
else:
out_channel = 1
if weight_conv != 'WeightConv2':
for i in range(num_layers):
self.masks.append(WeightConv1(hidden, hidden, out_channel))
else:
for i in range(num_layers):
self.masks.append(
WeightConv2(
Sequential(Linear(hidden * 2, hidden), ReLU(),
Linear(hidden, hidden), ReLU(),
Linear(hidden, out_channel), Sigmoid())))
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear(num_layers * hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
