| """ConvLSTM model for center-pixel next-day wildfire danger.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class ConvLSTMCell(nn.Module): |
| """Standard ConvLSTM cell with input, forget, output, and candidate gates.""" |
|
|
| def __init__(self, input_channels: int, hidden_channels: int, kernel_size: int = 3): |
| super().__init__() |
| self.hidden_channels = int(hidden_channels) |
| padding = kernel_size // 2 |
| self.gates = nn.Conv2d( |
| input_channels + hidden_channels, 4 * hidden_channels, |
| kernel_size=kernel_size, padding=padding, |
| ) |
|
|
| def forward(self, inputs: torch.Tensor, state: tuple[torch.Tensor, torch.Tensor]): |
| hidden, cell = state |
| input_gate, forget_gate, output_gate, candidate = self.gates( |
| torch.cat((inputs, hidden), dim=1) |
| ).chunk(4, dim=1) |
| input_gate = torch.sigmoid(input_gate) |
| forget_gate = torch.sigmoid(forget_gate) |
| output_gate = torch.sigmoid(output_gate) |
| candidate = torch.tanh(candidate) |
| next_cell = forget_gate * cell + input_gate * candidate |
| next_hidden = output_gate * torch.tanh(next_cell) |
| return next_hidden, next_cell |
|
|
|
|
| class FireCubeNet(nn.Module): |
| """Propagate ConvLSTM state over ten days and classify the center pixel.""" |
|
|
| def __init__(self, input_channels: int = 25, hidden_channels: int = 4, |
| kernel_size: int = 3, dropout: float = 0.1): |
| super().__init__() |
| self.input_channels = int(input_channels) |
| self.hidden_channels = int(hidden_channels) |
| self.cell = ConvLSTMCell(input_channels, hidden_channels, kernel_size) |
| self.head = nn.Sequential(nn.Dropout(dropout), nn.Linear(hidden_channels, 1)) |
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| if inputs.ndim != 5 or inputs.shape[2] != self.input_channels: |
| raise ValueError( |
| f"expected BTCHW with C={self.input_channels}, got {tuple(inputs.shape)}" |
| ) |
| batch, _, _, height, width = inputs.shape |
| hidden = inputs.new_zeros(batch, self.hidden_channels, height, width) |
| cell = inputs.new_zeros(batch, self.hidden_channels, height, width) |
| for time_index in range(inputs.shape[1]): |
| hidden, cell = self.cell(inputs[:, time_index], (hidden, cell)) |
| center_features = hidden[:, :, height // 2, width // 2] |
| return self.head(center_features) |
|
|
