import torch import numpy as np from torch import nn import torch.nn.functional as F from torch_geometric.nn import BatchNorm, global_add_pool, AttentiveFP, Set2Set from torch_geometric.nn.conv import MessagePassing from torch_geometric.nn.dense.linear import Linear as GeometricLinear from torch_geometric.utils import degree, scatter from torch.nn import ModuleList, Linear, ReLU, Sequential, Dropout class LegacyPNAConv(MessagePassing): """PyG 2.0.1 PNA behavior used to train the bundled checkpoint.""" def __init__(self, in_channels, out_channels, aggregators, scalers, deg, edge_dim=None, towers=1, pre_layers=1, post_layers=1, divide_input=False, **kwargs): kwargs.setdefault('aggr', None) super().__init__(node_dim=0, **kwargs) if divide_input: assert in_channels % towers == 0 assert out_channels % towers == 0 self.in_channels = in_channels self.out_channels = out_channels self.aggregators = aggregators self.scalers = scalers self.edge_dim = edge_dim self.towers = towers self.divide_input = divide_input self.F_in = in_channels // towers if divide_input else in_channels self.F_out = out_channels // towers deg = deg.to(torch.float) self.avg_deg = { 'lin': deg.mean().item(), 'log': (deg + 1).log().mean().item(), 'exp': deg.exp().mean().item(), } if edge_dim is not None: self.edge_encoder = GeometricLinear(edge_dim, self.F_in) self.pre_nns = ModuleList() self.post_nns = ModuleList() for _ in range(towers): modules = [GeometricLinear((3 if edge_dim else 2) * self.F_in, self.F_in)] for _ in range(pre_layers - 1): modules += [ReLU(), GeometricLinear(self.F_in, self.F_in)] self.pre_nns.append(Sequential(*modules)) post_in = (len(aggregators) * len(scalers) + 1) * self.F_in modules = [GeometricLinear(post_in, self.F_out)] for _ in range(post_layers - 1): modules += [ReLU(), GeometricLinear(self.F_out, self.F_out)] self.post_nns.append(Sequential(*modules)) self.lin = GeometricLinear(out_channels, out_channels) self.reset_parameters() def reset_parameters(self): if self.edge_dim is not None: self.edge_encoder.reset_parameters() for network in self.pre_nns: for module in network: if hasattr(module, 'reset_parameters'): module.reset_parameters() for network in self.post_nns: for module in network: if hasattr(module, 'reset_parameters'): module.reset_parameters() self.lin.reset_parameters() def forward(self, x, edge_index, edge_attr=None): if self.divide_input: x = x.view(-1, self.towers, self.F_in) else: x = x.view(-1, 1, self.F_in).repeat(1, self.towers, 1) out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=None) out = torch.cat([x, out], dim=-1) out = torch.cat([network(out[:, i]) for i, network in enumerate(self.post_nns)], dim=1) return self.lin(out) def message(self, x_i, x_j, edge_attr): if edge_attr is not None: edge_attr = self.edge_encoder(edge_attr) edge_attr = edge_attr.view(-1, 1, self.F_in).repeat(1, self.towers, 1) features = torch.cat([x_i, x_j, edge_attr], dim=-1) else: features = torch.cat([x_i, x_j], dim=-1) return torch.stack([network(features[:, i]) for i, network in enumerate(self.pre_nns)], dim=1) def aggregate(self, inputs, index, dim_size=None): outputs = [] for aggregator in self.aggregators: if aggregator in {'sum', 'mean', 'min', 'max'}: output = scatter(inputs, index, dim=0, dim_size=dim_size, reduce=aggregator) elif aggregator in {'var', 'std'}: mean = scatter(inputs, index, dim=0, dim_size=dim_size, reduce='mean') mean_squares = scatter(inputs * inputs, index, dim=0, dim_size=dim_size, reduce='mean') output = mean_squares - mean * mean if aggregator == 'std': output = torch.sqrt(torch.relu(output) + 1e-5) else: raise ValueError(f'Unknown aggregator "{aggregator}"') outputs.append(output) output = torch.cat(outputs, dim=-1) node_degree = degree(index, dim_size, dtype=inputs.dtype).clamp_(1).view(-1, 1, 1) outputs = [] for scaler in self.scalers: if scaler == 'identity': pass elif scaler == 'amplification': output = output * (torch.log(node_degree + 1) / self.avg_deg['log']) elif scaler == 'attenuation': output = output * (self.avg_deg['log'] / torch.log(node_degree + 1)) elif scaler == 'linear': output = output * (node_degree / self.avg_deg['lin']) elif scaler == 'inverse_linear': output = output * (self.avg_deg['lin'] / node_degree) else: raise ValueError(f'Unknown scaler "{scaler}"') outputs.append(output) return torch.cat(outputs, dim=-1) class MLP(nn.Module): def __init__(self,dims, n_layers, hidden_size, dropout=0 ): super().__init__() self.n_layers = n_layers self.hidden_size = hidden_size self.dims = dims self.dropout = dropout def block(in_size, n_hidden): layers = [ nn.Linear(in_size, n_hidden), nn.BatchNorm1d(n_hidden), nn.ReLU(), ] if self.dropout > 0: layers.append( nn.Dropout(self.dropout), ) return layers # Define PyTorch model self.model = nn.Sequential( *block(np.prod(dims), self.hidden_size) ) self.latent_size = self.hidden_size def forward(self, x): return self.model(x) class GNN_PNAConv(torch.nn.Module): def __init__(self, nodes_n_features: int, edges_n_features: int, deg: torch.tensor , n_layers: int = 6, hidden_size_node: int = 75, hidden_size_edges: int = 50, towers: int = 5, fcc_hidden_size: int = 50, dropout: float = 0, use_fds: bool = False, **args ): super(GNN_PNAConv, self).__init__() self.node_emb = MLP(nodes_n_features, 1, hidden_size_node) self.edge_emb = MLP(edges_n_features, 1, hidden_size_edges) aggregators = ['mean', 'min', 'max', 'std'] scalers = ['identity', 'amplification', 'attenuation'] self.convs = ModuleList() self.batch_norms = ModuleList() for _ in range(n_layers): conv = LegacyPNAConv(in_channels=hidden_size_node, out_channels=hidden_size_node, aggregators=aggregators, scalers=scalers, deg=deg, edge_dim=hidden_size_edges, towers=towers, pre_layers=2, post_layers=2, divide_input=False) self.convs.append(conv) self.batch_norms.append(BatchNorm(hidden_size_node)) self.set2set = Set2Set(hidden_size_node, processing_steps=6) fc_layers = [ Linear(2*hidden_size_node, hidden_size_node), BatchNorm(hidden_size_node), ReLU(), Linear(hidden_size_node, fcc_hidden_size), BatchNorm(fcc_hidden_size), ReLU(), ] self.fcc_hidden_size = fcc_hidden_size if dropout>0: fc_layers += [ Dropout(p=dropout) ] self.mlp = Sequential( *fc_layers ) def forward(self, x, edge_index, edge_attr, batch): x = self.node_emb(x.squeeze()) edge_attr = self.edge_emb(edge_attr) for conv, batch_norm in zip(self.convs, self.batch_norms): x = F.relu(batch_norm(conv(x, edge_index, edge_attr))) x = self.set2set(x, batch) #Set2Set #GlobalAttention #GraphMultisetTransformer return self.mlp(x) class QdolarAR(torch.nn.Module): def __init__(self, nodes_n_features: int, *args, **kargs ): super(QdolarAR, self).__init__() n_layers: int = 5 hidden_size_node: int = nodes_n_features fcc_hidden_size: int = 100 dropout: float = 0.28 self.fcc_hidden_size = fcc_hidden_size self.dropout = dropout self.n_layers = n_layers self.hidden_size_node = hidden_size_node assert n_layers >=3 fc_layers = [ BatchNorm(hidden_size_node), Linear( hidden_size_node, fcc_hidden_size), BatchNorm(fcc_hidden_size), ReLU(), ] for i in range(n_layers-2): fc_layers += [ Linear(fcc_hidden_size, fcc_hidden_size), BatchNorm(fcc_hidden_size), ReLU(), ] if dropout>0: fc_layers += [ Dropout(p=dropout) ] self.mlp = Sequential( *fc_layers ) def forward(self, x, edge_index, edge_attr, batch): assert len(batch) == len(batch.unique()) assert not torch.isnan(x).any(), "Error, x contains nan" return self.mlp(x) class GNN_AttentiveFP(torch.nn.Module): def __init__(self, nodes_n_features: int, edges_n_features: int, deg: torch.tensor , n_layers: int = 4, hidden_size_node: int = 75, hidden_size_edges: int = 50, towers: int = 5, fcc_hidden_size: int = 50, dropout: float = 0, use_fds: bool = False, **args ): super(GNN_AttentiveFP, self).__init__() self.attent_net = AttentiveFP(nodes_n_features, hidden_size_node, hidden_size_node, edges_n_features, n_layers, n_layers, dropout= dropout) fc_layers = [ Linear(hidden_size_node, hidden_size_edges), ReLU(), Linear(hidden_size_edges, fcc_hidden_size), ReLU(), ] self.fcc_hidden_size = fcc_hidden_size if dropout>0: fc_layers += Dropout(p=dropout) self.mlp = Sequential( *fc_layers ) def forward(self, x, edge_index, edge_attr, batch): x = self.attent_net(x, edge_index, edge_attr, batch) return self.mlp(x)