| import torch |
| import torch.nn as nn |
|
|
| from .edge_features import EdgeFeatures |
| from .self_attention import TransformerLayer, cat_neighbors_nodes, gather_nodes |
|
|
| class MVGNN(nn.Module): |
| def __init__(self, node_features, edge_features, hidden_dim, num_encoder_layers=4, k_neighbors=30, augment_eps=0., dropout=0.2): |
| super(MVGNN, self).__init__() |
|
|
| |
| self.augment_eps = augment_eps |
|
|
| |
| self.EdgeFeatures = EdgeFeatures(edge_features, top_k=k_neighbors, augment_eps=augment_eps) |
| self.dropout = nn.Dropout(dropout) |
|
|
| |
| self.act_fn = nn.ReLU() |
| self.lin1 = nn.Linear(node_features, 512, bias=True) |
| self.lin2 = nn.Linear(512, 256, bias=False) |
| self.W_v = nn.Linear(256, hidden_dim, bias=True) |
| self.W_e = nn.Linear(edge_features, hidden_dim, bias=True) |
| |
| |
| self.encoder_layers = nn.ModuleList([ |
| TransformerLayer(hidden_dim, hidden_dim * 2, dropout=dropout) |
| for _ in range(num_encoder_layers) |
| ]) |
| self.W_out = nn.Linear(hidden_dim * 2, hidden_dim, bias=True) |
| self.W_out1 = nn.Linear(hidden_dim, 64, bias=True) |
| self.W_out2 = nn.Linear(64, 1, bias=True) |
|
|
| |
| for p in self.parameters(): |
| if p.dim() > 1: |
| nn.init.xavier_uniform_(p) |
|
|
|
|
| def forward(self, X, V, mask,adj): |
| |
| |
| E, E_idx = self.EdgeFeatures(X, mask) |
|
|
| |
| if self.training and self.augment_eps > 0: |
| V = V + 0.1 * self.augment_eps * torch.randn_like(V) |
|
|
| V = self.act_fn(self.lin1(V)) |
| V = self.act_fn(self.lin2(V)) |
|
|
|
|
| h_V = self.W_v(V) |
| ho = h_V |
| h_E = self.W_e(E) |
|
|
| mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1) |
| mask_attend = mask.unsqueeze(-1) * mask_attend |
| for i,layer in enumerate(self.encoder_layers): |
| h_EV = cat_neighbors_nodes(h_V, h_E, E_idx) |
| h_V = layer(h_V, h_EV,adj,ho,i+1, mask_V=mask, mask_attend=mask_attend) |
| logits = self.act_fn(self.W_out1(h_V)) |
| logits = self.W_out2(logits).squeeze(-1) |
| |
| return logits |
|
|
