File size: 1,310 Bytes
cf5d356 | 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 | """Small neural-network layers used by RiboSphere."""
from __future__ import annotations
from torch import Tensor, nn
class FeedForward(nn.Module):
"""Two-layer MLP equivalent to the subset of timm.layers.Mlp used here."""
def __init__(
self,
input_features: int,
hidden_features: int,
output_features: int,
*,
activation: type[nn.Module] = nn.GELU,
dropout: float = 0.0,
) -> None:
super().__init__()
if min(input_features, hidden_features, output_features) <= 0:
raise ValueError("All feature dimensions must be positive.")
if not 0.0 <= dropout < 1.0:
raise ValueError("dropout must be in [0, 1).")
self.fc1 = nn.Linear(input_features, hidden_features)
self.activation = activation()
self.dropout1 = nn.Dropout(dropout)
self.fc2 = nn.Linear(hidden_features, output_features)
self.dropout2 = nn.Dropout(dropout)
def forward(self, inputs: Tensor) -> Tensor:
"""Apply the MLP without changing leading dimensions."""
outputs = self.fc1(inputs)
outputs = self.activation(outputs)
outputs = self.dropout1(outputs)
outputs = self.fc2(outputs)
outputs = self.dropout2(outputs)
return outputs
|