| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class LiquidTimeConstantRNN(nn.Module): |
| def __init__(self, hidden_dimensions: int = 41) -> None: |
| super().__init__() |
| self.hidden_dimensions = hidden_dimensions |
| self.candidate = nn.Linear(hidden_dimensions + 2, hidden_dimensions) |
| self.log_time_constant = nn.Parameter(torch.zeros(hidden_dimensions)) |
| self.output = nn.Linear(hidden_dimensions, 1) |
|
|
| def forward(self, sequence: torch.Tensor) -> torch.Tensor: |
| hidden = torch.zeros( |
| len(sequence), |
| self.hidden_dimensions, |
| device=sequence.device, |
| ) |
| time_constant = torch.nn.functional.softplus(self.log_time_constant) + 0.03 |
| outputs = [] |
| for step in range(sequence.shape[1]): |
| value_and_dt = sequence[:, step] |
| candidate_inputs = torch.cat( |
| [ |
| value_and_dt[:, :1], |
| value_and_dt[:, 1:2].clamp(max=0.12), |
| ], |
| dim=1, |
| ) |
| candidate = torch.tanh( |
| self.candidate(torch.cat([candidate_inputs, hidden], dim=1)) |
| ) |
| delta_time = value_and_dt[:, 1:2] |
| decay = torch.exp(-delta_time / time_constant[None]) |
| hidden = decay * hidden + (1 - decay) * candidate |
| outputs.append(self.output(hidden)) |
| return torch.stack(outputs, dim=1) |
|
|
|
|
| class MatchedGRU(nn.Module): |
| def __init__(self, hidden_dimensions: int = 23) -> None: |
| super().__init__() |
| self.recurrent = nn.GRU(2, hidden_dimensions, batch_first=True) |
| self.output = nn.Linear(hidden_dimensions, 1) |
|
|
| def forward(self, sequence: torch.Tensor) -> torch.Tensor: |
| hidden, _ = self.recurrent(sequence) |
| return self.output(hidden) |
|
|
|
|
| class MatchedRNN(nn.Module): |
| def __init__(self, hidden_dimensions: int = 41) -> None: |
| super().__init__() |
| self.recurrent = nn.RNN(2, hidden_dimensions, batch_first=True) |
| self.output = nn.Linear(hidden_dimensions, 1) |
|
|
| def forward(self, sequence: torch.Tensor) -> torch.Tensor: |
| hidden, _ = self.recurrent(sequence) |
| return self.output(hidden) |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
