File size: 5,216 Bytes
85851de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """
Self-contained molecular encoder for BioXMol.
This is the molecule branch (a Gated Graph Neural Network) of the BioXMol
soft-contrastive model. It maps a molecular graph to a fixed-length embedding.
No part of the mocop training framework is required to run it.
Architecture (as trained): n_edge=1, in_dim=75, n_conv=6, fc_dims=[1024, 128].
Embedding layers exposed by `embed(...)`:
"GNN" -> 75-d (graph-level sum-pooled node features)
"first_fc" -> 1024-d (after fc_layers[0]: Linear -> Tanh)
"second_fc" -> 128-d (after fc_layers[1]: Linear; contrastive projection)
"""
import math
from typing import Iterable, List
import torch
import torch.nn as nn
import torch.nn.functional as F
class GraphConvolution(nn.Module):
def __init__(self, in_dim: int, out_dim: int, n_edge: int = 1, bias: bool = True):
super().__init__()
self.weight = nn.Parameter(torch.Tensor(n_edge, in_dim, out_dim))
if bias:
self.bias = nn.Parameter(torch.Tensor(n_edge, 1, out_dim))
else:
self.register_parameter("bias", None)
self.reset_parameters()
def reset_parameters(self):
nn.init.xavier_uniform_(self.weight)
if self.bias is not None:
nn.init.zeros_(self.bias)
def forward(self, adj: torch.Tensor, feat: torch.Tensor) -> torch.Tensor:
if len(adj.size()) == 3:
adj = adj.unsqueeze(1)
feat = feat.unsqueeze(1)
output = torch.matmul(adj, feat)
output = torch.matmul(output, self.weight)
if self.bias is not None:
output = output + self.bias
output = output.sum(dim=1)
return F.relu(output)
class GRU2D(nn.Module):
"""2D GRU cell used to gate information flow between conv layers."""
def __init__(self, in_dim: int, hidden_dim: int, bias: bool = True):
super().__init__()
self.x_to_intermediate = nn.Linear(in_dim, 3 * hidden_dim, bias=bias)
self.h_to_intermediate = nn.Linear(in_dim, 3 * hidden_dim, bias=bias)
self.reset_parameters()
def reset_parameters(self):
for k, v in self.state_dict().items():
if "weight" in k:
std = math.sqrt(6.0 / (v.size(1) + v.size(0) / 3))
nn.init.uniform_(v, a=-std, b=std)
elif "bias" in k:
nn.init.zeros_(v)
def forward(self, x: torch.Tensor, h_0: torch.Tensor) -> torch.Tensor:
ix = self.x_to_intermediate(x)
ih = self.h_to_intermediate(h_0)
x_r, x_z, x_n = ix.chunk(3, -1)
h_r, h_z, h_n = ih.chunk(3, -1)
r = torch.sigmoid(x_r + h_r)
z = torch.sigmoid(x_z + h_z)
n = torch.tanh(x_n + (r * h_n))
return (1 - z) * n + z * h_0
class GatedGraphConvolution(nn.Module):
def __init__(self, in_dim: int, out_dim: int, n_edge: int = 1, bias: bool = True):
super().__init__()
if in_dim != out_dim:
raise ValueError(f"in_dim ({in_dim}) must equal out_dim ({out_dim}).")
self.gc = GraphConvolution(in_dim=in_dim, out_dim=out_dim, n_edge=n_edge, bias=bias)
self.gru = GRU2D(in_dim=out_dim, hidden_dim=out_dim, bias=bias)
def forward(self, adj: torch.Tensor, h_0: torch.Tensor) -> torch.Tensor:
h = self.gc(adj, h_0)
return self.gru(h, h_0)
class GatedGraphNeuralNetwork(nn.Module):
def __init__(
self,
n_edge: int = 1,
in_dim: int = 75,
n_conv: int = 6,
fc_dims: Iterable[int] = (1024, 128),
p_dropout: float = 0.0,
):
super().__init__()
fc_dims = list(fc_dims)
self.conv_layers = nn.ModuleList(
[GatedGraphConvolution(in_dim=in_dim, out_dim=in_dim, n_edge=n_edge) for _ in range(n_conv)]
)
num_fc = len(fc_dims)
dims = [in_dim] + fc_dims
fc_layers = []
for i, (a, b) in enumerate(zip(dims[:-1], dims[1:])):
layer = nn.Linear(a, b)
if i < num_fc - 2:
layer = nn.Sequential(layer, nn.ReLU())
elif i == num_fc - 2:
layer = nn.Sequential(layer, nn.Tanh())
fc_layers.append(layer)
self.fc_layers = nn.ModuleList(fc_layers)
self.dropout = nn.Dropout(p=p_dropout)
def _pool(self, adj, node_feat, atom_vec):
for conv in self.conv_layers:
node_feat = conv(adj, node_feat)
node_feat = self.dropout(node_feat)
return (node_feat * atom_vec).sum(1) # graph-level, masks padded atoms
def embed(self, adj, node_feat, atom_vec, layer: str = "first_fc") -> torch.Tensor:
"""Return embeddings from a chosen layer. Call model.eval() first."""
h = self._pool(adj, node_feat, atom_vec)
if layer == "GNN":
return h
h = self.fc_layers[0](h)
if layer == "first_fc":
return h
h = self.dropout(h)
h = self.fc_layers[1](h)
if layer == "second_fc":
return h
raise ValueError(f"Unknown layer '{layer}'. Use 'GNN', 'first_fc', or 'second_fc'.")
def forward(self, adj, node_feat, atom_vec):
return self.embed(adj, node_feat, atom_vec, layer="second_fc")
|