class MessagePassingAgnostic(torch.nn.Module):
"""
A model which does not perform any message passing.
Initial simplicial/cell representations are obtained by applying a dense layer, instead.
Sort of resembles a 'DeepSets'-likes architecture but on Simplicial/Cell Complexes.
"""
def __init__(self,
num_input_features,
num_classes,
hidden,
dropout_rate: float = 0.5,
max_dim: int = 2,
nonlinearity='relu',
readout='sum'):
super(MessagePassingAgnostic, self).__init__()
self.max_dim = max_dim
self.dropout_rate = dropout_rate
self.readout_type = readout
self.act = get_nonlinearity(nonlinearity, return_module=False)
self.lin0s = torch.nn.ModuleList()
for dim in range(max_dim + 1):
self.lin0s.append(Linear(num_input_features, hidden))
self.lin1 = Linear(hidden, hidden)
self.lin2 = Linear(hidden, num_classes)
def reset_parameters(self):
for lin0 in self.lin0s:
lin0.reset_parameters()
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data: ComplexBatch):
params = data.get_all_cochain_params(max_dim=self.max_dim,
include_down_features=False)
xs = list()
for dim in range(len(params)):
x_dim = params[dim].x
x_dim = self.lin0s[dim](x_dim)
xs.append(self.act(x_dim))
pooled_xs = pool_complex(xs, data, self.max_dim, self.readout_type)
pooled_xs = self.act(self.lin1(pooled_xs))
x = pooled_xs.sum(dim=0)
x = F.dropout(x, p=self.dropout_rate, training=self.training)
x = self.lin2(x)
return x
def __repr__(self):
return self.__class__.__name__
class ASAP(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden, ratio=0.8, dropout=0):
super().__init__()
self.conv1 = GraphConv(dataset.num_features, 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)
def reset_parameters(self):
self.conv1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
for pool in self.pools:
pool.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
edge_weight = None
x = F.relu(self.conv1(x, edge_index))
xs = [global_mean_pool(x, batch)]
for i, conv in enumerate(self.convs):
x = conv(x=x, edge_index=edge_index, edge_weight=edge_weight)
x = F.relu(x)
xs += [global_mean_pool(x, batch)]
if i % 2 == 0 and i < len(self.convs) - 1:
pool = self.pools[i // 2]
x, edge_index, edge_weight, batch, _ = pool(
x=x,
edge_index=edge_index,
edge_weight=edge_weight,
batch=batch)
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(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden, add_pool=False):
super(GIN, self).__init__()
self.conv1 = GINConv(Sequential(
Linear(dataset.num_features, hidden),
ReLU(),
Linear(hidden, hidden),
ReLU(),
BN(hidden),
),
train_eps=True)
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=True))
self.lin1 = Linear(hidden, hidden)
self.lin2 = Linear(hidden, dataset.num_classes)
self.add_pool = add_pool
def reset_parameters(self):
self.conv1.reset_parameters()
for conv in self.convs:
conv.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)
for conv in self.convs:
x = conv(x, edge_index)
if self.add_pool:
x = global_add_pool(x, batch)
else:
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 JKNet(torch.nn.Module):
def __init__(self,
in_channels,
hidden_channels,
out_channels,
num_layers,
dropout,
mode='cat'):
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=mode,
channels=hidden_channels,
num_layers=num_layers)
if mode == 'cat':
self.lin = Linear(num_layers * hidden_channels, out_channels)
else:
self.lin = 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.lin.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 = self.lin(x)
return F.log_softmax(x, dim=-1)
class GIN0WithJK(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden):
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='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 = 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 F.log_softmax(x, dim=-1)
def __repr__(self):
return self.__class__.__name__
class TopK(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden, ratio=0.25):
super(TopK, self).__init__()
self.conv1 = GNN_Block(dataset.num_features, hidden)
self.pool1 = TopKPooling(hidden, ratio)
self.convs = torch.nn.ModuleList()
self.pools = torch.nn.ModuleList()
self.convs.extend([
GNN_Block(hidden, hidden)
for i in range(num_layers - 1)
])
self.pools.extend(
[TopKPooling(hidden, ratio) for i in range((num_layers)-1)])
self.embed_final = GNN_Block(hidden, hidden)
self.jump = JumpingKnowledge(mode='cat')
self.lin1 = Linear((num_layers+1)*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()
for pool in self.pools:
pool.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 = [global_mean_pool(x, batch)]
x, edge_index, _, batch, _, _ = self.pool1(x, edge_index,
batch=batch)
for i, conv in enumerate(self.convs):
x = F.relu(conv(x, edge_index))
xs += [global_mean_pool(x, batch)]
pool = self.pools[i]
x, edge_index, _, batch, _, _ = pool(x, edge_index,
batch=batch)
x = self.embed_final(x, edge_index)
xs +=[global_mean_pool(x, batch)]
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 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.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.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, reg = dense_diff_pool(x, adj, s)
x = torch.cat(xs, dim=1)
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 InceptionBlock(torch.nn.Module):
def __init__(self, in_dim, out_dim):
super(InceptionBlock, self).__init__()
self.ln = Linear(in_dim, out_dim)
self.conv1 = DIGCNConv(in_dim, out_dim)
self.conv2 = DIGCNConv(in_dim, out_dim)
def reset_parameters(self):
self.ln.reset_parameters()
self.conv1.reset_parameters()
self.conv2.reset_parameters()
def forward(self, x, edge_index, edge_weight, edge_index2, edge_weight2):
x0 = self.ln(x)
x1 = self.conv1(x, edge_index, edge_weight)
x2 = self.conv2(x, edge_index2, edge_weight2)
return x0, x1, x2
class DefaultGraphHead(torch.nn.Module):
def __init__(self, hidden_channels, num_classes, dropout=0.5):
super().__init__()
self.lin1 = Linear(hidden_channels, hidden_channels)
self.lin2 = Linear(hidden_channels, num_classes)
self.dropout = dropout
def reset_parameters(self):
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, x):
x = F.relu(self.lin1(x))
x = F.dropout(x, p=self.dropout, training=self.training)
return self.lin2(x)
class GEN(torch.nn.Module):
def __init__(self, hidden_channels, num_layers):
super(GEN, self).__init__()
self.node_encoder = Linear(data.x.size(-1), hidden_channels)
self.edge_encoder = Linear(data.edge_attr.size(-1), hidden_channels)
self.layers = torch.nn.ModuleList()
for i in range(1, num_layers + 1):
conv = GENConv(hidden_channels,
hidden_channels,
aggr='softmax',
t=1.0,
learn_t=True,
num_layers=2,
norm='layer')
norm = LayerNorm(hidden_channels, elementwise_affine=True)
act = ReLU(inplace=True)
layer = DeepGCNLayer(conv,
norm,
act,
block='res+',
dropout=0.1,
ckpt_grad=i % 3)
self.layers.append(layer)
self.lin = Linear(hidden_channels, data.y.size(-1))
def reset_parameters(self):
self.node_encoder.reset_parameters()
self.edge_encoder.reset_parameters()
for layer in self.layers:
layer.reset_parameters()
def forward(self, x, edge_index, edge_attr):
x = self.node_encoder(x)
edge_attr = self.edge_encoder(edge_attr)
x = self.layers[0].conv(x, edge_index, edge_attr)
for layer in self.layers[1:]:
x = layer(x, edge_index, edge_attr)
x = self.layers[0].act(self.layers[0].norm(x))
x = F.dropout(x, p=0.1, training=self.training)
return self.lin(x)
class GatedGCN_directed(torch.nn.Module):
def __init__(self,num_layers=2,hidden=32,aggr='add',features_num=32,num_class=2,res=False,node_num=0):
super(GatedGCN_directed, self).__init__()
print(num_layers, aggr, hidden, res)
self.convs = torch.nn.ModuleList()
for i in range(num_layers):
self.convs.append(GatedBlock(hidden, aggr))
self.res = res
self.rnn = torch.nn.GRUCell(hidden * 2, hidden, bias=True)
self.first_lin = Linear(features_num, hidden)
self.fuse_weight = torch.nn.Parameter(torch.FloatTensor(num_layers),requires_grad=True)
self.fuse_weight.data.fill_(float(1) / (num_layers + 1))
self.rnn.reset_parameters()
self.out = Linear(hidden, num_class)
def reset_parameters(self):
self.first_lin.reset_parameters()
def forward(self, data):
x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight
edge_index2 = data.edge_index[[1, 0], :]
x = F.relu(self.first_lin(x))
x = F.dropout(x, p=0.5, training=self.training)
if (self.res == True):
i = 0
for conv in self.convs:
x_first = x
h = x
x1 = conv(x, edge_index, edge_weight=edge_weight)
x2 = conv(x, edge_index2, edge_weight=edge_weight)
x = torch.cat([x1, x2], dim=1)
x = self.rnn(x, h)
i += 1
else:
i = 0
for conv in self.convs:
x_first = x
h = x
x = conv(x, edge_index, edge_weight=edge_weight)
x = self.rnn(x, h)
if (i != 0):
x = x + self.fuse_weight[i] * x_first
i += 1
x = self.out(x)
return F.log_softmax(x, dim=-1)
def __repr__(self):
return self.__class__.__name__
class RevGNN(torch.nn.Module):
def __init__(self,
in_channels,
hidden_channels,
out_channels,
num_layers,
dropout,
num_groups=2):
super().__init__()
self.dropout = dropout
self.lin1 = Linear(in_channels, hidden_channels)
self.lin2 = Linear(hidden_channels, out_channels)
self.norm = LayerNorm(hidden_channels, elementwise_affine=True)
assert hidden_channels % num_groups == 0
self.convs = torch.nn.ModuleList()
for _ in range(num_layers):
conv = GNNBlock(
hidden_channels // num_groups,
hidden_channels // num_groups,
)
self.convs.append(GroupAddRev(conv, num_groups=num_groups))
def reset_parameters(self):
self.lin1.reset_parameters()
self.lin2.reset_parameters()
self.norm.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
def forward(self, x, edge_index):
x = self.lin1(x)
# Generate a dropout mask which will be shared across GNN blocks:
mask = None
if self.training and self.dropout > 0:
mask = torch.zeros_like(x).bernoulli_(1 - self.dropout)
mask = mask.requires_grad_(False)
mask = mask / (1 - self.dropout)
for conv in self.convs:
x = conv(x, edge_index, mask)
x = self.norm(x).relu()
x = F.dropout(x, p=self.dropout, training=self.training)
return self.lin2(x)
class GeoGraph(torch.nn.Module):
def __init__(self, hidden, geo, training):
super(GeoGraph, self).__init__()
self.geo = geo
self.training = training
if self.geo == True:
self.conv1 = GraphConv(64 + 3, hidden, aggr='mean')
else:
self.conv1 = GraphConv(64, hidden, aggr='mean')
self.conv2 = GraphConv(hidden, 32, aggr='mean')
self.conv3 = GraphConv(32, 16, aggr='mean')
self.lin1 = Linear(16, 16)
self.pool1 = EdgePoolingMod(16)
def reset_parameters(self):
self.conv1.reset_parameters()
self.conv2.reset_parameters()
self.conv3.reset_parameters()
self.lin1.reset_parameters()
def forward(self, data):
#Data(edge_index=[2, 210], neg_edge_index=[2, 182], pos_edge_index=[2, 28], x=[15, 512], x_bbox=[15, 4], x_heading=[15, 2], x_img_pos=[15, 2], x_pos=[15, 2], y=[210])
if self.training == True:
x, geo, reg, edge_index, edge_y = data.x.cuda(), data.geos.cuda(
), data.regressions.cuda(), data.edge_index.cuda(), data.y.cuda()
if self.geo == True:
# x = torch.cat((x,reg),1)
x = torch.cat((x, geo), 1)
else:
x, reg, edge_index = data.x.cuda(), data.regressions.cuda(
), data.edge_index.cuda()
x = F.relu(self.conv1(x, edge_index))
x = F.relu(self.conv2(x, edge_index))
x = F.relu(self.conv3(x, edge_index))
x = x.view(x.size()[0], -1)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.2, training=self.training)
x, edge_index, batch, edge_scores = self.pool1(x,
edge_index,
batch=None)
return edge_scores
def __repr__(self):
return self.__class__.__name__
class DenseJK(nn.Module):
def __init__(self, mode, channels=None, num_layers=None):
super(DenseJK, self).__init__()
self.channel = channels
self.mode = mode.lower()
assert self.mode in ['cat', 'max', 'lstm']
if mode == 'lstm':
assert channels is not None
assert num_layers is not None
self.lstm = LSTM(channels,
channels * num_layers // 2,
bidirectional=True,
batch_first=True)
self.att = Linear(2 * channels * num_layers // 2, 1)
self.reset_parameters()
def reset_parameters(self):
if hasattr(self, 'lstm'):
self.lstm.reset_parameters()
if hasattr(self, 'att'):
self.att.reset_parameters()
def forward(self, xs):
r"""Aggregates representations across different layers.
Args:
xs [batch, nodes, featdim*3]
"""
xs = torch.split(xs, self.channel, -1) # list of batch, node, featdim
xs = torch.stack(xs, 2) #[batch, nodes, num_layers, num_channels]
shape = xs.shape
x = xs.reshape(
(-1, shape[2],
shape[3])) # [ngraph * num_nodes , num_layers, num_channels]
alpha, _ = self.lstm(x)
alpha = self.att(alpha).squeeze(-1) # [ngraph * num_nodes, num_layers]
alpha = torch.softmax(alpha, dim=-1)
x = (x * alpha.unsqueeze(-1)).sum(dim=1)
x = x.reshape((shape[0], shape[1], shape[3]))
return x
def __repr__(self):
return '{}({})'.format(self.__class__.__name__, self.mode)
class DoubleNet(torch.nn.Module):
def __init__(self, dataset, num_features, num_classes):
super(DoubleNet, self).__init__()
self.conv1 = GCNConv(dataset.num_features, args.hidden)
self.conv2 = GCNConv(args.hidden, args.hidden)
self.conv1_ssl = GCNConv(dataset.num_features, args.hidden)
self.conv2_ssl = GCNConv(args.hidden, int(dataset[0].num_class))
self.lin = Linear(int(dataset[0].num_class), int(dataset[0].num_class))
def reset_parameters(self):
self.conv1.reset_parameters()
self.conv2.reset_parameters()
self.conv1_ssl.reset_parameters()
self.conv2_ssl.reset_parameters()
self.lin.reset_parameters()
def decoder(self, z, edge_index, sigmoid=True):
value = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=1)
return torch.sigmoid(value) if sigmoid else value
def forward(self, data, pos_edge_index, neg_edge_index, edge_index,
masked_nodes):
x = data.x
x = F.relu(self.conv1(x, edge_index)) # LAYER 1
z = self.conv2(x, edge_index) # LAYER 2
total_edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1)
pos_pred = self.decoder(z, pos_edge_index, sigmoid=True)
neg_pred = self.decoder(z, neg_edge_index, sigmoid=True)
total_pred = torch.cat([pos_pred, neg_pred], dim=-1)
r, c = total_edge_index[0][total_pred > 0.5], total_edge_index[1][
total_pred > 0.5]
new_index = torch.stack(
(torch.cat([r, c], dim=-1), (torch.cat([c, r], dim=-1))), dim=0)
added_index = torch.cat([edge_index, new_index], dim=-1)
x = data.x
x = F.relu(self.conv1_ssl(x, added_index)) # LAYER 1
x = self.conv2_ssl(x, added_index) # LAYER 2
# return self.lin(F.relu(x)), z,r.size()
out = self.lin(F.relu(x))
drop = torch.nn.Dropout(p=0.5)
return F.log_softmax(x, dim=1), z, out, r.size(-1)
class GCN(torch.nn.Module):
def __init__(self, num_layers=2, hidden=16, features_num=16, num_class=2):
super(GCN, self).__init__()
# first layer
self.conv1 = GCNConv(features_num, hidden)
# list of 2nd - num_layers layers
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(GCNConv(hidden, hidden))
# fully connected layers
self.lin2 = Linear(hidden, num_class)
self.first_lin = Linear(features_num, hidden)
def reset_parameters(self):
# clear weights
self.first_lin.reset_parameters()
self.conv1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight
# fully connected layer + relu
x = F.relu(self.first_lin(x))
# dropout layer
x = F.dropout(x, p=0.5, training=self.training)
# GCN layers
for conv in self.convs:
x = F.relu(conv(x, edge_index, edge_weight=edge_weight))
# Another dropout
x = F.dropout(x, p=0.5, training=self.training)
# second FC layer
x = self.lin2(x)
# Softmax
return F.log_softmax(x, dim=-1)
def __repr__(self):
return self.__class__.__name__
class Net(torch.nn.Module):
def __init__(self, dataset):
super(Net, self).__init__()
self.conv1 = GCNConv(dataset.num_features, hidden)
# self.conv2 = GCNConv(hidden, int(dataset[0].num_class))
self.conv2 = GCNConv(dataset.num_features, int(dataset[0].num_class))
self.lin = Linear(int(dataset[0].num_class),int(dataset[0].num_class))
def reset_parameters(self):
self.conv1.reset_parameters()
self.conv2.reset_parameters()
self.lin.reset_parameters()
def forward(self, data, edge_index):
x= data.x
# x = F.relu(self.conv1(x, edge_index))
x = self.conv2(x, edge_index)
return self.lin(F.relu(x))
class GraphSAGEWithJK(torch.nn.Module):
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)
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 = torch.cat([
global_add_pool(x, batch),
global_mean_pool(x, batch),
global_max_pool(x, batch)
],
dim=1)
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(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, num_conv_layers, dropout=0.5):
super(GIN, self).__init__()
self.convs = torch.nn.ModuleList()
self.convs.append(
GINConv(Linear(in_channels, hidden_channels), train_eps=True)
)
for _ in range(num_conv_layers-1):
self.convs.append(
GINConv(Linear(hidden_channels, hidden_channels), train_eps=True)
)
self.pooling = global_mean_pool
self.classify = Linear(hidden_channels, out_channels)
self.dropout = dropout
def reset_parameters(self):
for conv in self.convs:
conv.reset_parameters()
self.classify.reset_parameters()
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = x.float()
# convs -> node embedding
for i, conv in enumerate(self.convs):
x = conv(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
# pooling -> graph embedding
x = self.pooling(x, data.batch)
# linear -> classification
x = self.classify(x)
return x
class SGCN(torch.nn.Module):
def __init__(self,
num_layers=2,
hidden=16,
features_num=16,
num_class=2,
hidden_droprate=0.5,
edge_droprate=0.0):
super(SGCN, self).__init__()
self.conv1 = SGConv(features_num, hidden)
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(SGConv(hidden, hidden))
self.lin2 = Linear(hidden, num_class)
self.first_lin = Linear(features_num, hidden)
self.hidden_droprate = hidden_droprate
self.edge_droprate = edge_droprate
def reset_parameters(self):
self.first_lin.reset_parameters()
self.conv1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
if self.edge_droprate != 0.0:
x = data.x
edge_index, edge_weight = dropout_adj(data.edge_index,
data.edge_weight,
self.edge_droprate)
else:
x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight
x = F.relu(self.first_lin(x))
x = F.dropout(x, p=self.dropout_rate, training=self.training)
for conv in self.convs:
x = F.relu(conv(x, edge_index, edge_weight=edge_weight))
x = F.dropout(x, p=self.dropout_rate, training=self.training)
x = self.lin2(x)
# return F.log_softmax(x, dim=-1)
# due to focal loss: return the logits, put the log_softmax operation into the GNNAlgo
return x
def __repr__(self):
return self.__class__.__name__
class Net(torch.nn.Module):
def __init__(self):
super().__init__()
hidden = args.hidden
num_layers = 5
ratio = 0.8
self.conv1 = GraphConv(dataset.num_features, hidden, aggr='add')
self.convs = torch.nn.ModuleList()
self.pools = torch.nn.ModuleList()
self.convs.extend([
GraphConv(hidden, hidden, aggr='add')
for i in range(num_layers - 1)
])
self.pools.extend(
[TopKPooling(hidden, ratio) 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)
def reset_parameters(self):
self.conv1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
for pool in self.pools:
pool.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 = [global_add_pool(x, batch)]
for i, conv in enumerate(self.convs):
x = F.relu(conv(x, edge_index))
xs += [global_add_pool(x, batch)]
if i % 2 == 0 and i < len(self.convs) - 1:
pool = self.pools[i // 2]
x, edge_index, _, batch, _, _ = pool(x,
edge_index,
batch=batch)
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)
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.lin1 = Linear(data.x.shape[1], 64)
self.lin2 = Linear(64, int(max(data.y)) + 1)
self.prop1 = APPNP(10, 0.1)
def reset_parameters(self):
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, x, edge_index):
x = F.dropout(x, p=0.5, training=self.training)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
x = self.prop1(x, edge_index)
return F.log_softmax(x, dim=1)
class APPNP_Net(torch.nn.Module):
def __init__(self, input_channels, hidden_channels, out_channels, K=10, alpha=0.1):
super().__init__()
self.lin1 = Linear(input_channels, hidden_channels)
self.lin2 = Linear(hidden_channels, out_channels)
self.prop = APPNP(K, alpha)
def reset_parameters(self):
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, x, edge_index, dropout_ratio=0.6):
x = F.dropout(x, p=dropout_ratio, training=self.training)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=dropout_ratio, training=self.training)
x = self.lin2(x)
x = self.prop(x, edge_index)
return F.log_softmax(x, dim=1)
class GCNWithJK(torch.nn.Module):
def __init__(self,
num_features,
output_channels,
num_layers=3,
nb_neurons=128,
mode='cat',
**kwargs):
super(GCNWithJK, self).__init__()
self.conv1 = GCNConv(num_features, nb_neurons)
self.convs = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(GCNConv(nb_neurons, nb_neurons))
self.jump = JumpingKnowledge(mode)
if mode == 'cat':
self.lin1 = Linear(num_layers * nb_neurons, nb_neurons)
else:
self.lin1 = Linear(nb_neurons, nb_neurons)
self.lin2 = Linear(nb_neurons, output_channels)
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, target_size, **kwargs):
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, size=target_size)
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 ASAP_Pool(torch.nn.Module):
def __init__(self, dataset, num_layers, hidden, ratio=0.8, **kwargs):
super(ASAP_Pool, self).__init__()
if type(ratio) != list:
ratio = [ratio for i in range(num_layers)]
self.conv1 = GCNConv(dataset.num_features, hidden)
self.pool1 = ASAP_Pooling(in_channels=hidden, ratio=ratio[0], **kwargs)
self.convs = torch.nn.ModuleList()
self.pools = torch.nn.ModuleList()
for i in range(num_layers - 1):
self.convs.append(GCNConv(hidden, hidden))
self.pools.append(ASAP_Pooling(in_channels=hidden, ratio=ratio[i], **kwargs))
self.lin1 = Linear(2 * hidden, hidden) # 2*hidden due to readout layer
self.lin2 = Linear(hidden, dataset.num_classes)
self.reset_parameters()
def reset_parameters(self):
self.conv1.reset_parameters()
self.pool1.reset_parameters()
for conv, pool in zip(self.convs, self.pools):
conv.reset_parameters()
pool.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))
x, edge_index, edge_weight, batch, perm = self.pool1(x=x, edge_index=edge_index, edge_weight=None, batch=batch)
xs = readout(x, batch)
for conv, pool in zip(self.convs, self.pools):
x = F.relu(conv(x=x, edge_index=edge_index, edge_weight=edge_weight))
x, edge_index, edge_weight, batch, perm = pool(x=x, edge_index=edge_index, edge_weight=edge_weight,
batch=batch)
xs += readout(x, batch)
x = F.relu(self.lin1(xs))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
out = F.log_softmax(x, dim=-1)
return out
def __repr__(self):
return self.__class__.__name__
class Net(torch.nn.Module):
def __init__(self, dataset):
super(Net, self).__init__()
self.lin1 = Linear(dataset.num_features, args.hidden)
self.lin2 = Linear(args.hidden, dataset.num_classes)
self.prop1 = APPNP(args.K, args.alpha)
def reset_parameters(self):
self.lin1.reset_parameters()
self.lin2.reset_parameters()
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = F.dropout(x, p=args.dropout, training=self.training)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=args.dropout, training=self.training)
x = self.lin2(x)
x = self.prop1(x, edge_index)
return F.log_softmax(x, dim=1)
class GNN(torch.nn.Module):
# a base GNN class, GCN message passing + sum_pooling
def __init__(self, dataset, gconv=GCNConv, latent_dim=[32, 32, 32, 1], regression=False, adj_dropout=0.2, force_undirected=False):
super(GNN, self).__init__()
self.regression = regression
self.adj_dropout = adj_dropout
self.force_undirected = force_undirected
self.convs = torch.nn.ModuleList()
self.convs.append(gconv(dataset.num_features, latent_dim[0]))
for i in range(0, len(latent_dim)-1):
self.convs.append(gconv(latent_dim[i], latent_dim[i+1]))
self.lin1 = Linear(sum(latent_dim), 128)
if self.regression:
self.lin2 = Linear(128, 1)
else:
self.lin2 = Linear(128, dataset.num_classes)
def reset_parameters(self):
for conv in self.convs:
conv.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
if self.adj_dropout > 0:
edge_index, edge_type = dropout_adj(edge_index, edge_type, p=self.adj_dropout, force_undirected=self.force_undirected, num_nodes=len(x), training=self.training)
concat_states = []
for conv in self.convs:
x = torch.tanh(conv(x, edge_index))
concat_states.append(x)
concat_states = torch.cat(concat_states, 1)
x = global_add_pool(concat_states, batch)
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.lin2(x)
if self.regression:
return x[:, 0]
else:
return F.log_softmax(x, dim=-1)
def __repr__(self):
return self.__class__.__name__
class GNN_Block(torch.nn.Module):
def __init__(self, in_channels, hidden_channels):
super(GNN_Block, self).__init__()
self.conv1 = GCNConv(in_channels, hidden_channels)
self.conv2 = GCNConv(hidden_channels, hidden_channels)
self.lin = Linear(hidden_channels + hidden_channels, hidden_channels)
def reset_parameters(self):
self.conv1.reset_parameters()
self.conv2.reset_parameters()
self.lin.reset_parameters()
def forward(self, x, edge_index):
x1 = F.relu(self.conv1(x, edge_index))
x2 = F.relu(self.conv2(x1, edge_index))
out = self.lin(torch.cat((x1, x2), -1))
return out
class ModelGCN(torch.nn.Module):
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 reset_parameters(self):
self.linear_1.reset_parameters()
for conv in self.convs:
conv.reset_parameters()
self.linear_2.reset_parameters()
def forward(self, data):
x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight
edge_index, edge_weight = dropout_adj(edge_index,
edge_attr=edge_weight,
p=0.8,
training=self.training)
x_jk = []
x = self.linear_1(x)
x = self.activation(x)
x_jk.append(dropout(x, p=0.5, training=self.training))
for i in range(len(self.convs)):
x = self.convs[i](x_jk[-1], edge_index, edge_weight=edge_weight)
if i != len(self.convs) - 1:
x_jk.append(self.activation(x))
else:
x_jk.append(dropout(x, p=0.5, training=self.training))
x = self.JK(x_jk)
x = self.linear_2(x)
return log_softmax(x, dim=-1)
class Block(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels):
super(Block, self).__init__()
self.conv1 = DenseGCNConv(in_channels, hidden_channels)
self.conv2 = DenseGCNConv(hidden_channels, hidden_channels)
self.lin = Linear(hidden_channels + hidden_channels, out_channels)
def reset_parameters(self):
self.conv1.reset_parameters()
self.conv2.reset_parameters()
self.lin.reset_parameters()
def forward(self, x, adj, mask=None, add_loop=True):
x1 = F.relu(self.conv1(x, adj, mask, add_loop))
x2 = F.relu(self.conv2(x1, adj, mask, add_loop))
out = self.lin(torch.cat((x1, x2), -1))
return out
