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