| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class ClockworkRNN(nn.Module): |
| def __init__(self, hidden_dimensions: int = 32) -> None: |
| super().__init__() |
| self.hidden_dimensions = hidden_dimensions |
| self.periods = [1, 2, 4, 8] |
| self.block_size = hidden_dimensions // len(self.periods) |
| self.input_projection = nn.Linear(1, hidden_dimensions) |
| self.recurrent_projection = nn.Linear(hidden_dimensions, hidden_dimensions) |
| self.output = nn.Linear(hidden_dimensions, 1) |
|
|
| def forward(self, sequence: torch.Tensor) -> torch.Tensor: |
| batch, length, _ = sequence.shape |
| hidden = torch.zeros(batch, self.hidden_dimensions, device=sequence.device) |
| outputs = [] |
| for step in range(length): |
| candidate = torch.tanh( |
| self.input_projection(sequence[:, step]) |
| + self.recurrent_projection(hidden) |
| ) |
| blocks = [] |
| for index, period in enumerate(self.periods): |
| start = index * self.block_size |
| end = start + self.block_size |
| blocks.append( |
| candidate[:, start:end] if step % period == 0 else hidden[:, start:end] |
| ) |
| hidden = torch.cat(blocks, dim=1) |
| outputs.append(self.output(hidden)) |
| return torch.stack(outputs, dim=1) |
|
|
|
|
| class PlainRNN(nn.Module): |
| def __init__(self, hidden_dimensions: int = 32) -> None: |
| super().__init__() |
| self.recurrent = nn.RNN(1, 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 MatchedGRU(nn.Module): |
| def __init__(self, hidden_dimensions: int = 18) -> None: |
| super().__init__() |
| self.recurrent = nn.GRU(1, 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()) |
|
|
|
|
