File size: 5,776 Bytes
8f4ed7a | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """Neural network architectures for descriptor and fingerprint baselines."""
from __future__ import annotations
from typing import List, Optional, Tuple, cast
import torch
import torch.nn as nn
class DescriptorNN(nn.Module):
"""Feed-forward network for molecular descriptors with optional lab embedding."""
def __init__(
self,
input_dim: int,
hidden_dims: List[int] | Tuple[int, ...] = (192, 96, 48),
dropout: float = 0.2,
use_batch_norm: bool = True,
num_labs: Optional[int] = None,
lab_embed_dim: int = 12,
input_dropout: float = 0.05,
) -> None:
super().__init__()
self.use_lab_embedding = num_labs is not None and num_labs > 0
lab_count = cast(int, num_labs) if self.use_lab_embedding else 0
self.lab_embedding = (
nn.Embedding(lab_count, lab_embed_dim)
if self.use_lab_embedding
else None
)
initial_dim = input_dim + (lab_embed_dim if self.use_lab_embedding else 0)
self.input_dropout = nn.Dropout(input_dropout) if input_dropout > 0 else None
layers: List[nn.Module] = []
prev_dim = initial_dim
for hidden_dim in hidden_dims:
layers.append(nn.Linear(prev_dim, hidden_dim))
if use_batch_norm:
layers.append(nn.BatchNorm1d(hidden_dim))
layers.append(nn.GELU())
layers.append(nn.Dropout(dropout))
prev_dim = hidden_dim
bottleneck_dim = max(32, prev_dim // 2)
layers.extend(
[
nn.Linear(prev_dim, bottleneck_dim),
nn.GELU(),
nn.Dropout(dropout * 0.5),
nn.Linear(bottleneck_dim, 1),
]
)
self.network = nn.Sequential(*layers)
self.apply(self._init_weights)
# placeholders populated during training for inverse scaling
self.target_mean: float = 0.0
self.target_std: float = 1.0
@staticmethod
def _init_weights(module: nn.Module) -> None:
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.BatchNorm1d):
nn.init.ones_(module.weight)
nn.init.zeros_(module.bias)
def forward(self, x: torch.Tensor, lab_indices: Optional[torch.Tensor] = None) -> torch.Tensor:
if self.input_dropout is not None:
x = self.input_dropout(x)
if self.lab_embedding is not None:
if lab_indices is None:
raise ValueError("lab_indices must be provided when lab embedding is enabled.")
if lab_indices.dim() > 1:
lab_indices = lab_indices.squeeze(-1)
lab_embeddings = self.lab_embedding(lab_indices.long())
x = torch.cat([x, lab_embeddings], dim=-1)
return self.network(x).squeeze(-1)
class FingerprintNN(nn.Module):
"""Feed-forward network tailored for high-dimensional sparse fingerprints."""
def __init__(
self,
input_dim: int,
hidden_dims: List[int] | Tuple[int, ...] = (768, 384, 192, 96),
dropout: float = 0.15,
use_batch_norm: bool = True,
num_labs: Optional[int] = None,
lab_embed_dim: int = 16,
input_dropout: float = 0.1,
) -> None:
super().__init__()
self.use_lab_embedding = num_labs is not None and num_labs > 0
lab_count = cast(int, num_labs) if self.use_lab_embedding else 0
self.lab_embedding = (
nn.Embedding(lab_count, lab_embed_dim)
if self.use_lab_embedding
else None
)
initial_dim = input_dim + (lab_embed_dim if self.use_lab_embedding else 0)
self.input_dropout = nn.Dropout(input_dropout) if input_dropout > 0 else None
layers: List[nn.Module] = []
prev_dim = initial_dim
for idx, hidden_dim in enumerate(hidden_dims):
layers.append(nn.Linear(prev_dim, hidden_dim))
if use_batch_norm:
layers.append(nn.BatchNorm1d(hidden_dim))
layers.append(nn.GELU())
layers.append(nn.Dropout(dropout if idx < len(hidden_dims) - 1 else dropout * 0.5))
prev_dim = hidden_dim
layers.append(nn.Linear(prev_dim, 1))
self.network = nn.Sequential(*layers)
self.apply(self._init_weights)
self.target_mean: float = 0.0
self.target_std: float = 1.0
@staticmethod
def _init_weights(module: nn.Module) -> None:
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.BatchNorm1d):
nn.init.ones_(module.weight)
nn.init.zeros_(module.bias)
def forward(self, x: torch.Tensor, lab_indices: Optional[torch.Tensor] = None) -> torch.Tensor:
if self.input_dropout is not None:
x = self.input_dropout(x)
if self.lab_embedding is not None:
if lab_indices is None:
raise ValueError("lab_indices must be provided when lab embedding is enabled.")
if lab_indices.dim() > 1:
lab_indices = lab_indices.squeeze(-1)
lab_embeddings = self.lab_embedding(lab_indices.long())
x = torch.cat([x, lab_embeddings], dim=-1)
return self.network(x).squeeze(-1)
__all__ = ["DescriptorNN", "FingerprintNN"] |