import torch.nn.functional as F import torch from torch_geometric.nn import GCNConv, GATv2Conv, SAGEConv class GCNConvModel(torch.nn.Module): """ Two layers of GCNConv. Matches SOTA performance with 20% parameters. """ def __init__(self, in_channels=1024, feature_channels=128, additional_channels=20, out_channels=3, dropout=0.7, use_additional_channels=False): super(GCNConvModel, self).__init__() self.use_additional_channels = use_additional_channels if self.use_additional_channels: in_channels = in_channels + additional_channels self.gcn1 = GCNConv(in_channels=in_channels, out_channels=feature_channels) self.gcn2 = GCNConv(in_channels=feature_channels, out_channels=out_channels) self.dropout = dropout def forward(self, features, edges, edges2, edge_features, additional_feature): """Second pair of edge tensor allows different shells of neighbors to be evaluated. Not utilized in this model, but data pipeline must work for all models.""" if self.use_additional_channels: features = torch.concat([features, additional_feature], dim=1) edge_features = None x = self.gcn1(features, edges2, edge_features) x = F.relu(x) x = F.dropout(x, self.dropout, training=self.training) return self.gcn2(x, edges2, edge_features) class SAGEConvModel(torch.nn.Module): """ Two layers of SAGEConv. """ def __init__(self, in_channels=1024, feature_channels=128, additional_channels=20, out_channels=3, dropout=0.7, activation=F.leaky_relu, use_additional_channels=False): super(SAGEConvModel, self).__init__() self.use_additional_channels = use_additional_channels if self.use_additional_channels: in_channels = in_channels + additional_channels self.sage1 = SAGEConv(in_channels=in_channels, out_channels=feature_channels, aggr="mean") self.sage2 = SAGEConv(in_channels=feature_channels, out_channels=out_channels, aggr="mean") self.dropout = dropout self.activation = activation def forward(self, features, edges, edges2, edge_features, additional_feature): if self.use_additional_channels: features = torch.concat([features, additional_feature], dim=1) x = self.sage1(features, edges2) x = self.activation(x) x = F.dropout(x, self.dropout, training=self.training) return self.sage2(x, edges2) class SAGEConvMLPModel(torch.nn.Module): """ SAGEConv followed by MLP. This model is numerically strongest. """ def __init__(self, in_channels=1024, feature_channels=128, additional_channels=20, out_channels=3, dropout=0.7, activation=F.leaky_relu, heads=4, dropout_fcn=0.5, use_additional_channels=False): super(SAGEConvMLPModel, self).__init__() self.use_additional_channels = use_additional_channels if self.use_additional_channels: intermediate_layer_size = feature_channels + additional_channels else: intermediate_layer_size = feature_channels self.sage1 = SAGEConv(in_channels=in_channels, out_channels=feature_channels, aggr="mean") # Why // head? What is heads? self.fc1 = torch.nn.Linear(intermediate_layer_size, intermediate_layer_size // heads) self.fc2 = torch.nn.Linear(intermediate_layer_size // heads, out_channels) self.norm = torch.nn.BatchNorm1d(intermediate_layer_size // heads) self.dropout = dropout self.dropout_fcn = dropout_fcn self.activation = activation def forward(self, features, edges, edges2, edge_features, additional_feature): x = self.sage1(features, edges2) x = self.activation(x) x = F.dropout(x, self.dropout, training=self.training) # if DSSP features are used they are combined with the features extracted from the embeddings by the SAGEConv # layer; numerically, it performed better to feed in DSSP feats here if self.use_additional_channels: x = self.fc1(torch.concat([x, additional_feature], dim=1)) else: x = self.fc1(x) x = F.relu(x) x = self.norm(x) x = F.dropout(x, self.dropout_fcn, training=self.training) return self.fc2(x) class SAGEConvGATMLPModel(torch.nn.Module): """ Mix of SAGEConv, GAT and MLP model with different neighbor shells evaluated. """ def __init__(self, in_channels=1024, feature_channels=128, additional_channels=20, out_channels=3, dropout=0.7, activation=F.leaky_relu, heads=4, dropout_fcn=0.5, use_additional_channels=False): super(SAGEConvGATMLPModel, self).__init__() self.use_additional_channels = use_additional_channels if self.use_additional_channels: sage2_channels = additional_channels else: sage2_channels = in_channels self.sage1 = SAGEConv(in_channels=in_channels, out_channels=feature_channels, aggr="mean") self.sage2 = SAGEConv(in_channels=sage2_channels, out_channels=feature_channels, aggr="mean") self.gat1 = GATv2Conv(in_channels=feature_channels, out_channels=int(feature_channels / 2), heads=heads, edge_dim=1, concat=True) self.gat2 = GATv2Conv(in_channels=heads * int(feature_channels / 2), out_channels=int(feature_channels / 4), heads=heads, edge_dim=1, concat=True) self.fc1 = torch.nn.Linear( feature_channels + int(feature_channels / 4) * heads, feature_channels + int(feature_channels / 4) * heads ) self.fc2 = torch.nn.Linear( feature_channels + int(feature_channels / 4) * heads, out_channels ) self.normx = torch.nn.BatchNorm1d(feature_channels) self.normy = torch.nn.BatchNorm1d(int(feature_channels / 4) * heads) self.normfcn = torch.nn.BatchNorm1d(feature_channels + int(feature_channels / 4) * heads) self.dropout = dropout self.dropout_fcn = dropout_fcn self.activation = activation def forward(self, features, edges, edges2, edge_features, additional_features): if self.use_additional_channels: gat_input = self.normx(self.sage2(additional_features, edges)) else: gat_input = self.normx(self.sage2(features, edges)) x = self.sage1(features, edges) x = self.normx(x) y = self.gat1(gat_input, edges2, edge_features) y = self.activation(y) y = F.dropout(y, self.dropout, training=self.training) y = self.gat2(y, edges2, edge_features) y = self.normy(y) z = torch.concat([x, y], dim=1) z = self.activation(z) z = F.dropout(z, self.dropout, training=self.training) z = self.fc1(z) z = self.normfcn(z) z = F.relu(z) z = F.dropout(z, self.dropout_fcn, training=self.training) return self.fc2(z)