| |
| """UnifiedIMUNet -- dilated 1-D ResNet with attention pooling. |
| |
| Architecture ported verbatim from sample_code/ts-tcmpio-unified-temporal-resnet.ipynb |
| (cell 15). Single shared trunk, one velocity head; the platform head is an |
| auxiliary training-time signal only and is not used at inference, so the model |
| stays a single unified network with no external routing. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| def make_norm(kind: str, channels: int) -> nn.Module: |
| """'bn' BatchNorm1d; 'gn' GroupNorm with 8 groups -- statistics per sample, |
| so a mixed-platform batch no longer normalises a hovering drone's features |
| against a running car's.""" |
| if kind == "bn": |
| return nn.BatchNorm1d(channels) |
| if kind == "gn": |
| return nn.GroupNorm(min(8, channels), channels) |
| raise ValueError(f"norm must be bn or gn, got {kind!r}") |
|
|
|
|
| class ResidualBlock(nn.Module): |
| def __init__(self, in_channels, out_channels, stride=1, dilation=1, norm="bn"): |
| super().__init__() |
| padding = 2 * dilation |
| self.main = nn.Sequential( |
| nn.Conv1d(in_channels, out_channels, 5, stride=stride, padding=padding, |
| dilation=dilation, bias=False), |
| make_norm(norm, out_channels), |
| nn.GELU(), |
| nn.Conv1d(out_channels, out_channels, 5, padding=padding, dilation=dilation, bias=False), |
| make_norm(norm, out_channels), |
| ) |
| self.skip = (nn.Conv1d(in_channels, out_channels, 1, stride=stride, bias=False) |
| if stride != 1 or in_channels != out_channels else nn.Identity()) |
| self.activation = nn.GELU() |
|
|
| def forward(self, x): |
| return self.activation(self.main(x) + self.skip(x)) |
|
|
|
|
| class UnifiedIMUNet(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.encoder = nn.Sequential( |
| nn.Conv1d(6, 64, 7, padding=3, bias=False), nn.BatchNorm1d(64), nn.GELU(), |
| ResidualBlock(64, 64), |
| ResidualBlock(64, 96, stride=2), |
| ResidualBlock(96, 96, dilation=2), |
| ResidualBlock(96, 128, stride=2, dilation=2), |
| ResidualBlock(128, 128, dilation=4), |
| ResidualBlock(128, 192, stride=2, dilation=4), |
| ResidualBlock(192, 192, dilation=4), |
| ) |
| self.attention = nn.Conv1d(192, 1, 1) |
| self.shared = nn.Sequential(nn.LayerNorm(576), nn.Linear(576, 256), nn.GELU(), nn.Dropout(0.2)) |
| self.velocity_head = nn.Linear(256, 3) |
| self.platform_head = nn.Linear(256, 4) |
|
|
| def features_and_heads(self, x): |
| """(N, 6, 200) -> (feature (N, 256), velocity (N, 3), platform_logits (N, 4)). |
| |
| Auxiliary heads need the shared feature, which forward() discards. |
| """ |
| x = self.encoder(x) |
| weights = torch.softmax(self.attention(x), dim=-1) |
| features = torch.cat([(x * weights).sum(-1), x.mean(-1), x.std(-1, unbiased=False)], dim=1) |
| features = self.shared(features) |
| return features, self.velocity_head(features), self.platform_head(features) |
|
|
| def forward(self, x): |
| _, velocity, platform_logits = self.features_and_heads(x) |
| return velocity, platform_logits |
|
|
|
|
| def vector_huber(prediction, target, beta=0.25): |
| error = torch.linalg.vector_norm(prediction - target, dim=1) |
| return torch.where(error < beta, 0.5 * error.square() / beta, error - 0.5 * beta).mean() |
|
|
