| """Compact U-Net architecture for the released WildFIRE-FM checkpoints.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def make_norm(norm_type: str, num_channels: int, norm_groups: int) -> nn.Module: |
| if norm_type == "batch": |
| return nn.BatchNorm2d(num_channels) |
| if norm_type == "group": |
| groups = max(1, min(int(norm_groups), num_channels)) |
| while num_channels % groups != 0 and groups > 1: |
| groups -= 1 |
| return nn.GroupNorm(groups, num_channels) |
| if norm_type == "instance": |
| return nn.InstanceNorm2d(num_channels, affine=True) |
| if norm_type in {"none", "identity"}: |
| return nn.Identity() |
| raise ValueError(f"Unsupported norm_type: {norm_type}") |
|
|
|
|
| class ConvBlock(nn.Module): |
| def __init__(self, in_ch: int, out_ch: int, norm_type: str, norm_groups: int): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), |
| make_norm(norm_type, out_ch, norm_groups), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), |
| make_norm(norm_type, out_ch, norm_groups), |
| nn.ReLU(inplace=True), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.net(x) |
|
|
|
|
| class UNetSmallFlex(nn.Module): |
| def __init__( |
| self, |
| in_ch: int, |
| base: int = 32, |
| dropout: float = 0.1, |
| norm_type: str = "group", |
| norm_groups: int = 8, |
| prior_prob: float | None = None, |
| use_aux_spatial_head: bool = False, |
| aux_prior_prob: float | None = None, |
| ): |
| super().__init__() |
| self.enc1 = ConvBlock(in_ch, base, norm_type, norm_groups) |
| self.enc2 = ConvBlock(base, base * 2, norm_type, norm_groups) |
| self.enc3 = ConvBlock(base * 2, base * 4, norm_type, norm_groups) |
| self.enc4 = ConvBlock(base * 4, base * 8, norm_type, norm_groups) |
| self.pool = nn.MaxPool2d(2) |
| self.bottleneck = ConvBlock(base * 8, base * 16, norm_type, norm_groups) |
| self.up4 = nn.ConvTranspose2d(base * 16, base * 8, 2, stride=2) |
| self.dec4 = ConvBlock(base * 16, base * 8, norm_type, norm_groups) |
| self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2) |
| self.dec3 = ConvBlock(base * 8, base * 4, norm_type, norm_groups) |
| self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2) |
| self.dec2 = ConvBlock(base * 4, base * 2, norm_type, norm_groups) |
| self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2) |
| self.dec1 = ConvBlock(base * 2, base, norm_type, norm_groups) |
| self.drop = nn.Dropout2d(p=dropout) |
| self.head = nn.Conv2d(base, 1, kernel_size=1) |
| self.use_aux_spatial_head = bool(use_aux_spatial_head) |
| self.aux_head = nn.Conv2d(base, 1, kernel_size=1) if self.use_aux_spatial_head else None |
| if prior_prob is not None: |
| prior_prob = float(min(max(prior_prob, 1e-6), 1.0 - 1e-6)) |
| nn.init.constant_(self.head.bias, math.log(prior_prob / (1.0 - prior_prob))) |
| if self.aux_head is not None and aux_prior_prob is not None: |
| aux_prior_prob = float(min(max(aux_prior_prob, 1e-6), 1.0 - 1e-6)) |
| nn.init.constant_(self.aux_head.bias, math.log(aux_prior_prob / (1.0 - aux_prior_prob))) |
|
|
| @staticmethod |
| def _match_hw(x: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: |
| diff_y = ref.size(2) - x.size(2) |
| diff_x = ref.size(3) - x.size(3) |
| if diff_y > 0 or diff_x > 0: |
| x = F.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) |
| if diff_y < 0: |
| y0 = (-diff_y) // 2 |
| x = x[:, :, y0 : y0 + ref.size(2), :] |
| if diff_x < 0: |
| x0 = (-diff_x) // 2 |
| x = x[:, :, :, x0 : x0 + ref.size(3)] |
| return x |
|
|
| def forward(self, x: torch.Tensor, return_aux: bool = False): |
| e1 = self.enc1(x) |
| e2 = self.enc2(self.pool(e1)) |
| e3 = self.enc3(self.pool(e2)) |
| e4 = self.enc4(self.pool(e3)) |
| b = self.bottleneck(self.pool(e4)) |
| d4 = self.dec4(torch.cat([self._match_hw(self.up4(b), e4), e4], dim=1)) |
| d3 = self.dec3(torch.cat([self._match_hw(self.up3(d4), e3), e3], dim=1)) |
| d2 = self.dec2(torch.cat([self._match_hw(self.up2(d3), e2), e2], dim=1)) |
| d1 = self.dec1(torch.cat([self._match_hw(self.up1(d2), e1), e1], dim=1)) |
| features = self.drop(d1) |
| logits = self.head(features) |
| if return_aux and self.aux_head is not None: |
| return logits, self.aux_head(features) |
| return logits |
|
|
