|
|
|
|
|
|
|
|
|
|
|
from argparse import Namespace |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
from .features import GVPGraphEmbedding |
|
from .gvp_modules import GVPConvLayer, LayerNorm |
|
from .gvp_utils import unflatten_graph |
|
|
|
|
|
|
|
class GVPEncoder(nn.Module): |
|
|
|
def __init__(self, args): |
|
super().__init__() |
|
self.args = args |
|
self.embed_graph = GVPGraphEmbedding(args) |
|
|
|
node_hidden_dim = (args.node_hidden_dim_scalar, |
|
args.node_hidden_dim_vector) |
|
edge_hidden_dim = (args.edge_hidden_dim_scalar, |
|
args.edge_hidden_dim_vector) |
|
|
|
conv_activations = (F.relu, torch.sigmoid) |
|
self.encoder_layers = nn.ModuleList( |
|
GVPConvLayer( |
|
node_hidden_dim, |
|
edge_hidden_dim, |
|
drop_rate=args.dropout, |
|
vector_gate=True, |
|
attention_heads=0, |
|
n_message=3, |
|
conv_activations=conv_activations, |
|
n_edge_gvps=0, |
|
eps=1e-4, |
|
layernorm=True, |
|
) |
|
for i in range(args.num_encoder_layers) |
|
) |
|
|
|
def forward(self, coords, coord_mask, padding_mask, confidence): |
|
node_embeddings, edge_embeddings, edge_index = self.embed_graph( |
|
coords, coord_mask, padding_mask, confidence) |
|
|
|
for i, layer in enumerate(self.encoder_layers): |
|
node_embeddings, edge_embeddings = layer(node_embeddings, |
|
edge_index, edge_embeddings) |
|
|
|
node_embeddings = unflatten_graph(node_embeddings, coords.shape[0]) |
|
return node_embeddings |
|
|
