| """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)
|
|
|
|
|
| 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"] |
