From 5f5de6ebee60b795b6de9bb21347468a996be8d8 Mon Sep 17 00:00:00 2001 From: Denys Senkin Date: Thu, 22 Jan 2026 19:16:07 +0100 Subject: [PATCH] Add group transform module --- timesead_ext/models/__init__.py | 3 + timesead_ext/models/transforms/__init__.py | 3 + timesead_ext/models/transforms/group.py | 136 +++++++++++++++++++++ 3 files changed, 142 insertions(+) create mode 100644 timesead_ext/models/transforms/__init__.py create mode 100644 timesead_ext/models/transforms/group.py diff --git a/timesead_ext/models/__init__.py b/timesead_ext/models/__init__.py index d84011f..2cdfe35 100644 --- a/timesead_ext/models/__init__.py +++ b/timesead_ext/models/__init__.py @@ -1 +1,4 @@ from . import other +from . import transforms + +__all__ = ["other", "transforms"] diff --git a/timesead_ext/models/transforms/__init__.py b/timesead_ext/models/transforms/__init__.py new file mode 100644 index 0000000..43cf6da --- /dev/null +++ b/timesead_ext/models/transforms/__init__.py @@ -0,0 +1,3 @@ +from .group import GroupTransform + +__all__ = ["GroupTransform"] diff --git a/timesead_ext/models/transforms/group.py b/timesead_ext/models/transforms/group.py new file mode 100644 index 0000000..5901db0 --- /dev/null +++ b/timesead_ext/models/transforms/group.py @@ -0,0 +1,136 @@ +"""Learnable group-wise temporal transforms.""" + +from __future__ import annotations + +from typing import Dict, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _softplus_inverse(value: torch.Tensor) -> torch.Tensor: + return torch.log(torch.expm1(value)) + + +def _make_base_grid(length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + return torch.linspace(-1.0, 1.0, length, device=device, dtype=dtype).unsqueeze(0) + + +def _grid_sample_1d(x: torch.Tensor, grid_t: torch.Tensor) -> torch.Tensor: + x_2d = x.unsqueeze(-1) + n, _, length = x.shape + grid = torch.zeros((n, length, 1, 2), device=x.device, dtype=x.dtype) + grid[..., 1] = grid_t.unsqueeze(-1) + return F.grid_sample(x_2d, grid, mode="bilinear", padding_mode="border", align_corners=True).squeeze(-1) + + +def _apply_scale_bias( + x: torch.Tensor, + scale: Optional[torch.Tensor], + bias: Optional[torch.Tensor], +) -> torch.Tensor: + if scale is not None: + x = x * scale + if bias is not None: + x = x + bias + return x + + +class GroupTransform(nn.Module): + """Apply differentiable group-wise time shift, scale/bias, and monotone warp.""" + + def __init__(self, group_cfg: Dict): + super().__init__() + self.k_group = int(group_cfg.get("k_group", 1)) + self.enable_shift = bool(group_cfg.get("enable_shift", True)) + self.enable_scale_bias = bool(group_cfg.get("enable_scale_bias", True)) + self.enable_warp = bool(group_cfg.get("enable_warp", True)) + + self.warp_knots = int(group_cfg.get("warp_knots", 8)) + self.max_shift_frac = float(group_cfg.get("max_shift_frac", 0.0)) + self.scale_range = float(group_cfg.get("scale_range", 0.0)) + self.bias_range = float(group_cfg.get("bias_range", 0.0)) + self.warp_strength = float(group_cfg.get("warp_strength", 0.0)) + + if self.enable_shift: + self.shift = nn.Parameter(torch.zeros(self.k_group)) + else: + self.register_parameter("shift", None) + + if self.enable_scale_bias: + init_scale = _softplus_inverse(torch.tensor(1.0)) + self.scale_raw = nn.Parameter(init_scale.repeat(self.k_group)) + self.bias_raw = nn.Parameter(torch.zeros(self.k_group)) + else: + self.register_parameter("scale_raw", None) + self.register_parameter("bias_raw", None) + + if self.enable_warp: + init_knots = _softplus_inverse(torch.ones(self.warp_knots)) + self.warp_raw = nn.Parameter(init_knots.repeat(self.k_group, 1)) + else: + self.register_parameter("warp_raw", None) + + def _expand_group_params(self, params: torch.Tensor, batch_size: int) -> torch.Tensor: + return params.repeat(batch_size) + + def _get_scale_bias(self, batch_size: int, channels: int, length: int) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + if not self.enable_scale_bias: + return None, None + + scale = F.softplus(self.scale_raw) + scale = 1.0 + self.scale_range * (scale - 1.0) + bias = self.bias_range * torch.tanh(self.bias_raw) + + scale = self._expand_group_params(scale, batch_size).view(batch_size * self.k_group, 1, 1) + bias = self._expand_group_params(bias, batch_size).view(batch_size * self.k_group, 1, 1) + scale = scale.expand(-1, channels, length) + bias = bias.expand(-1, channels, length) + return scale, bias + + def _get_shift_grid(self, batch_size: int, length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + base = _make_base_grid(length, device, dtype) + shift = self._expand_group_params(self.shift, batch_size) + shift = torch.tanh(shift) * self.max_shift_frac + shift = shift * 2.0 + return base + shift.unsqueeze(-1) + + def _get_warp_grid(self, batch_size: int, length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + knots = F.softplus(self.warp_raw) + knots = torch.cumsum(knots, dim=-1) + knots = knots / knots[..., -1:].clamp_min(1e-6) + + knots = knots.unsqueeze(1) + warp = F.interpolate(knots, size=length, mode="linear", align_corners=True).squeeze(1) + + identity = torch.linspace(0.0, 1.0, length, device=device, dtype=dtype).unsqueeze(0) + warp = identity + self.warp_strength * (warp - identity) + warp = warp.clamp(0.0, 1.0) + warp = self._expand_group_params(warp, batch_size) + return warp.view(batch_size * self.k_group, length) * 2.0 - 1.0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: Input tensor of shape (batch, channels, length). + + Returns: + Tensor of shape (batch, k_group, channels, length). + """ + batch_size, channels, length = x.shape + x_grouped = x.unsqueeze(1).expand(batch_size, self.k_group, channels, length) + x_grouped = x_grouped.reshape(batch_size * self.k_group, channels, length) + + if self.enable_shift and self.max_shift_frac > 0: + grid_t = self._get_shift_grid(batch_size, length, x.device, x.dtype) + x_grouped = _grid_sample_1d(x_grouped, grid_t) + + if self.enable_warp and self.warp_strength > 0: + grid_t = self._get_warp_grid(batch_size, length, x.device, x.dtype) + x_grouped = _grid_sample_1d(x_grouped, grid_t) + + scale, bias = self._get_scale_bias(batch_size, channels, length) + x_grouped = _apply_scale_bias(x_grouped, scale, bias) + + return x_grouped.view(batch_size, self.k_group, channels, length)