Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import dgl.function as fn | |
| class EGNNConv(nn.Module): | |
| def __init__(self, in_size, hidden_size, out_size, edge_feat_size=0): | |
| super(EGNNConv, self).__init__() | |
| self.in_size = in_size | |
| self.hidden_size = hidden_size | |
| self.out_size = out_size | |
| self.edge_feat_size = edge_feat_size | |
| act_fn = nn.SiLU() | |
| # \phi_e | |
| self.edge_mlp = nn.Sequential( | |
| # +1 for the radial feature: ||x_i - x_j||^2 | |
| nn.Linear(in_size * 2 + edge_feat_size + 1, hidden_size), | |
| act_fn, | |
| nn.Linear(hidden_size, hidden_size), | |
| act_fn | |
| ) | |
| # \phi_h | |
| self.node_mlp = nn.Sequential( | |
| nn.Linear(in_size + hidden_size, hidden_size), | |
| act_fn, | |
| nn.Linear(hidden_size, out_size) | |
| ) | |
| # \phi_x | |
| self.coord_mlp = nn.Sequential( | |
| nn.Linear(hidden_size, hidden_size), | |
| act_fn, | |
| nn.Linear(hidden_size, 1, bias=False) | |
| ) | |
| def message(self, edges): | |
| """message function for EGNN""" | |
| # concat features for edge mlp | |
| if self.edge_feat_size > 0: | |
| f = torch.cat( | |
| [edges.src['h'], edges.dst['h'], edges.data['radial'], edges.data['a']], | |
| dim=-1 | |
| ) | |
| else: | |
| f = torch.cat([edges.src['h'], edges.dst['h'], edges.data['radial']], dim=-1) | |
| msg_h = self.edge_mlp(f) | |
| msg_x = self.coord_mlp(msg_h) * edges.data['x_diff'] | |
| return {'msg_x': msg_x, 'msg_h': msg_h} | |
| def forward(self, graph, node_feat, coord_feat, edge_feat=None): | |
| with graph.local_scope(): | |
| # node feature | |
| graph.ndata['h'] = node_feat | |
| # coordinate feature | |
| graph.ndata['x'] = coord_feat | |
| # edge feature | |
| if self.edge_feat_size > 0: | |
| assert edge_feat is not None, "Edge features must be provided." | |
| graph.edata['a'] = edge_feat | |
| # get coordinate diff & radial features | |
| graph.apply_edges(fn.u_sub_v('x', 'x', 'x_diff')) | |
| graph.edata['radial'] = graph.edata['x_diff'].square().sum(dim=1).unsqueeze(-1) | |
| # normalize coordinate difference | |
| graph.edata['x_diff'] = graph.edata['x_diff'] / (graph.edata['radial'].sqrt() + 1e-30) | |
| graph.apply_edges(self.message) | |
| graph.update_all(fn.copy_e('msg_x', 'm'), fn.mean('m', 'x_neigh')) | |
| graph.update_all(fn.copy_e('msg_h', 'm'), fn.sum('m', 'h_neigh')) | |
| h_neigh, x_neigh = graph.ndata['h_neigh'], graph.ndata['x_neigh'] | |
| h = self.node_mlp( | |
| torch.cat([node_feat, h_neigh], dim=-1) | |
| ) | |
| x = coord_feat + x_neigh | |
| return h, x | |
| class EGNN(nn.Module): | |
| def __init__(self, input_node_dim, input_edge_dim, hidden_dim, num_layers, dropout, JK='sum'): | |
| super(EGNN, self).__init__() | |
| self.num_layers = num_layers | |
| # List of MLPs | |
| self.egnn_layers = torch.nn.ModuleList() | |
| self.batch_norms = torch.nn.ModuleList() | |
| for layer in range(self.num_layers - 1): | |
| if layer == 0: | |
| self.egnn_layers.append(EGNNConv(input_node_dim, hidden_dim, hidden_dim, input_edge_dim)) | |
| else: | |
| self.egnn_layers.append(EGNNConv(hidden_dim, hidden_dim, hidden_dim, input_edge_dim)) | |
| self.batch_norms.append(nn.BatchNorm1d(hidden_dim)) | |
| self.drop = nn.Dropout(dropout) | |
| self.JK = JK | |
| def forward(self, g, Perturb=None): | |
| hidden_rep = [] | |
| node_feats = g.ndata.pop('h').float() | |
| edge_feats = g.edata['e'] | |
| coord_feats = g.ndata['pos'] | |
| for idx, egnn in enumerate(self.egnn_layers): | |
| if idx == 0 and Perturb is not None: | |
| node_feats = node_feats + Perturb | |
| node_feats, coord_feats = egnn(g, node_feats, coord_feats, edge_feats) | |
| node_feats = self.batch_norms[idx](node_feats) | |
| node_feats = F.relu(node_feats) | |
| node_feats = self.drop(node_feats) | |
| hidden_rep.append(node_feats) | |
| if self.JK == 'sum': | |
| hidden_rep = [h.unsqueeze(0) for h in hidden_rep] | |
| return torch.sum(torch.cat(hidden_rep, dim=0), dim=0) | |
| elif self.JK == 'max': | |
| hidden_rep = [h.unsqueeze(0) for h in hidden_rep] | |
| return torch.max(torch.cat(hidden_rep, dim=0), dim=0)[0] | |
| elif self.JK == 'concat': | |
| return torch.cat(hidden_rep, dim=1) | |
| elif self.JK == 'last': | |
| return hidden_rep[-1] | |