File size: 3,448 Bytes
1b0d082 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | #!/usr/bin/env python3
"""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()
|