File size: 1,775 Bytes
4483e82 | 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 | from __future__ import annotations
import torch
from torch import nn
from .numpy_runtime import DEFAULT_CENTER, DEFAULT_SCALE, WINDOW_SIZE
class TelemetryFeatureExtractor(nn.Module):
"""Export-friendly implementation of the six physical window features."""
def __init__(self) -> None:
super().__init__()
x = torch.arange(WINDOW_SIZE, dtype=torch.float32)
x_centered = x - x.mean()
self.register_buffer("x_centered", x_centered)
self.register_buffer("slope_denominator", torch.square(x_centered).sum())
def forward(self, telemetry: torch.Tensor) -> torch.Tensor:
force = telemetry[:, :, 0]
deviation = telemetry[:, :, 1]
slope = (force * self.x_centered).sum(dim=1) / self.slope_denominator
shift = force[:, -10:].mean(dim=1) - force[:, :10].mean(dim=1)
std = force.std(dim=1, correction=0)
max_deviation = deviation.amax(dim=1)
last_deviation = deviation[:, -1]
force_range = force.amax(dim=1) - force.amin(dim=1)
return torch.stack(
[slope, shift, std, max_deviation, last_deviation, force_range],
dim=1,
)
class TinyDriftNet(nn.Module):
def __init__(self) -> None:
super().__init__()
self.features = TelemetryFeatureExtractor()
self.register_buffer("feature_center", torch.tensor(DEFAULT_CENTER.copy()))
self.register_buffer("feature_scale", torch.tensor(DEFAULT_SCALE.copy()))
self.classifier = nn.Linear(6, 1)
def forward(self, telemetry: torch.Tensor) -> torch.Tensor:
features = self.features(telemetry)
normalized = (features - self.feature_center) / self.feature_scale
return torch.sigmoid(self.classifier(normalized)).squeeze(1)
|