Upload modeling_imu_encoder.py with huggingface_hub

modeling_imu_encoder.py

@@ -0,0 +1,317 @@
+"""WISDM IMU Masked Encoder — self-supervised Transformer for activity recognition.
+
+A pure-PyTorch model for encoding 10-second IMU sensor windows (6-channel
+accelerometer + gyroscope @ 20 Hz) into 192-dim representations. Pretrained
+with masked prediction, SupCon contrastive learning, and LMM frequency loss
+on the WISDM smartphone+smartwatch dataset (18 activity classes).
+
+Usage:
+    from modeling_imu_encoder import IMUMaskedEncoder
+
+    model = IMUMaskedEncoder.from_pretrained("NikoKKK/IMU-SelfSupEncoder-v1")
+    model.eval()
+
+    # Input: (B, 6, 200) tensor — 6 IMU channels, 200 timesteps
+    with torch.no_grad():
+        patch_out, intermediates, cls_out, global_freq = model(x)
+    # cls_out: (B, 192) — global representation for classification
+    # patch_out: (B, 20, 192) — per-window-patch embeddings
+    # intermediates: {layer_idx: (B, 20, 192)} — intermediate layer outputs
+"""
+
+import math
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+class IMUEncoderConfig:
+    """Configuration for IMUMaskedEncoder.
+
+    Attributes:
+        n_channels (int): Number of IMU sensor channels (default: 6).
+        patch_size (int): Conv-stem patch size in timesteps (default: 10).
+        n_patches (int): Number of patches per window (default: 20).
+        embed_dim (int): Embedding dimension (default: 192).
+        n_layers (int): Number of Transformer encoder layers (default: 4).
+        n_heads (int): Number of attention heads (default: 6).
+        mlp_ratio (float): MLP hidden/embed_dim ratio (default: 3.0).
+        dropout (float): Dropout probability (default: 0.1).
+        target_layers (List[int]): Layers collected as intermediate outputs.
+    """
+
+    def __init__(
+        self,
+        n_channels: int = 6,
+        patch_size: int = 10,
+        n_patches: int = 20,
+        embed_dim: int = 192,
+        n_layers: int = 4,
+        n_heads: int = 6,
+        mlp_ratio: float = 3.0,
+        dropout: float = 0.1,
+        target_layers: Optional[List[int]] = None,
+        **kwargs,
+    ):
+        self.n_channels = n_channels
+        self.patch_size = patch_size
+        self.n_patches = n_patches
+        self.embed_dim = embed_dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+        self.target_layers = target_layers or [2, 4]
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "IMUEncoderConfig":
+        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})
+
+    def to_dict(self) -> dict:
+        return {
+            "n_channels": self.n_channels,
+            "patch_size": self.patch_size,
+            "n_patches": self.n_patches,
+            "embed_dim": self.embed_dim,
+            "n_layers": self.n_layers,
+            "n_heads": self.n_heads,
+            "mlp_ratio": self.mlp_ratio,
+            "dropout": self.dropout,
+            "target_layers": self.target_layers,
+        }
+
+
+# ---------------------------------------------------------------------------
+# Sub-modules
+# ---------------------------------------------------------------------------
+
+class PatchTimeFreqEmbedding(nn.Module):
+    """Conv-stem per-patch time + local frequency fusion -> embed_dim tokens."""
+
+    def __init__(self, config: IMUEncoderConfig):
+        super().__init__()
+        half_dim = config.embed_dim // 2
+        self.conv_stem = nn.Conv1d(
+            config.n_channels, half_dim, kernel_size=config.patch_size,
+            stride=config.patch_size, bias=False,
+        )
+        n_freq_bins = config.patch_size // 2 + 1
+        self.freq_proj = nn.Linear(config.n_channels * n_freq_bins, half_dim)
+        self.pos_embed = nn.Parameter(
+            torch.randn(1, config.n_patches + 1, config.embed_dim) * 0.02
+        )
+        self.fusion = nn.Linear(config.embed_dim, config.embed_dim)
+        self.norm = nn.LayerNorm(config.embed_dim)
+        self.patch_size = config.patch_size
+        self.n_patches = config.n_patches
+
+        # FFT over full 200-timestep signal: rfft(200) → 101 bins, then .mean(dim=1)
+        n_global_freq_bins = config.patch_size * config.n_patches // 2 + 1
+        self.global_freq_proj = nn.Linear(n_global_freq_bins, config.embed_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        B, C, T = x.shape
+        time_feat = self.conv_stem(x)
+
+        x_patch = x.view(B, C, self.n_patches, self.patch_size)
+        x_patch_centered = x_patch - x_patch.mean(dim=-1, keepdim=True)
+        freq_mag = torch.abs(torch.fft.rfft(x_patch_centered, dim=-1))
+        n_freq = freq_mag.size(-1)
+        freq_flat = freq_mag.permute(0, 2, 1, 3).reshape(B, self.n_patches, C * n_freq)
+        freq_feat = self.freq_proj(freq_flat).transpose(1, 2)
+
+        combined = torch.cat([time_feat, freq_feat], dim=1)
+        tokens = self.fusion(combined.transpose(1, 2))
+
+        full_freq = torch.abs(torch.fft.rfft(x, dim=-1)).mean(dim=1)
+        global_freq = self.global_freq_proj(full_freq).unsqueeze(1)
+
+        tokens = tokens + self.pos_embed[:, :self.n_patches, :]
+        global_freq = global_freq + self.pos_embed[:, self.n_patches:self.n_patches + 1, :]
+        tokens = self.norm(tokens)
+        global_freq = self.norm(global_freq)
+        return tokens, global_freq
+
+
+class TransformerEncoder(nn.Module):
+    """Standard Transformer encoder, collects intermediate layer outputs."""
+
+    def __init__(self, config: IMUEncoderConfig):
+        super().__init__()
+        self.n_layers = config.n_layers
+        self.collect_layers = config.target_layers
+        self.cls_token = nn.Parameter(torch.randn(1, 1, config.embed_dim) * 0.02)
+        self.layers = nn.ModuleList([
+            nn.TransformerEncoderLayer(
+                d_model=config.embed_dim,
+                nhead=config.n_heads,
+                dim_feedforward=int(config.embed_dim * config.mlp_ratio),
+                dropout=config.dropout,
+                activation="gelu",
+                batch_first=True,
+                norm_first=True,
+            )
+            for _ in range(config.n_layers)
+        ])
+        self.final_norm = nn.LayerNorm(config.embed_dim)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x: torch.Tensor):
+        B = x.shape[0]
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat([cls_tokens, self.dropout(x)], dim=1)
+
+        intermediates = {}
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+            layer_idx = i + 1
+            if layer_idx in self.collect_layers:
+                intermediates[layer_idx] = self.final_norm(x[:, 1:, :])
+
+        x = self.final_norm(x)
+        cls_out = x[:, 0, :]
+        patch_out = x[:, 1:, :]
+        if self.n_layers in self.collect_layers:
+            intermediates[self.n_layers] = patch_out
+        return patch_out, intermediates, cls_out
+
+
+# ---------------------------------------------------------------------------
+# Main model
+# ---------------------------------------------------------------------------
+
+class IMUMaskedEncoder(nn.Module):
+    """Self-supervised encoder for IMU-based human activity recognition.
+
+    Input:  (B, 6, 200) — 6-channel IMU window (accel_x/y/z, gyro_x/y/z) @ 20 Hz.
+    Output: patch_out (B, 20, 192), intermediates dict, cls_out (B, 192),
+            global_freq (B, 1, 192).
+
+    The CLS token (`cls_out`) is the primary global representation for
+    downstream classification. Intermediate layer outputs can be used for
+    multi-level feature extraction or distillation.
+
+    Usage:
+        model = IMUMaskedEncoder.from_pretrained("NikoKKK/IMU-SelfSupEncoder-v1")
+        model.eval()
+
+        x = torch.randn(8, 6, 200)  # 8 windows of 10 seconds each
+        with torch.no_grad():
+            patch_out, intermediates, cls_out, global_freq = model(x)
+
+        print(cls_out.shape)  # (8, 192)
+    """
+
+    _HUB_URL = "https://huggingface.co/NikoKKK/IMU-SelfSupEncoder-v1"
+
+    def __init__(self, config: IMUEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.embed = PatchTimeFreqEmbedding(config)
+        self.transformer = TransformerEncoder(config)
+        self.mask_token = nn.Parameter(torch.randn(1, 1, config.embed_dim) * 0.02)
+
+    @classmethod
+    def from_pretrained(cls, model_id: str = "NikoKKK/IMU-SelfSupEncoder-v1",
+                        force_download: bool = False) -> "IMUMaskedEncoder":
+        """Load pretrained model from Hugging Face Hub or local path.
+
+        Args:
+            model_id: Hugging Face model ID (e.g. "NikoKKK/IMU-SelfSupEncoder-v1")
+                      or local directory path.
+            force_download: Force re-download from Hub.
+
+        Returns:
+            IMUMaskedEncoder with pretrained weights loaded.
+        """
+        from huggingface_hub import hf_hub_download
+        from safetensors.torch import load_file
+
+        # Determine if model_id is a local path
+        if os.path.isdir(model_id):
+            config_path = os.path.join(model_id, "config.json")
+            weights_path = os.path.join(model_id, "model.safetensors")
+        else:
+            config_path = hf_hub_download(
+                model_id, "config.json", force_download=force_download,
+            )
+            weights_path = hf_hub_download(
+                model_id, "model.safetensors", force_download=force_download,
+            )
+
+        with open(config_path, "r") as f:
+            config_dict = json.load(f)
+        config = IMUEncoderConfig.from_dict(config_dict)
+
+        model = cls(config)
+        state_dict = load_file(weights_path)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def forward(self, x: torch.Tensor, mask_matrix=None):
+        """Forward pass.
+
+        Args:
+            x: (B, C, T) tensor — IMU sensor window.
+               Expected C=6 (accel_x,y,z + gyro_x,y,z), T=200 @ 20 Hz.
+            mask_matrix: Optional (B, N_patches) boolean tensor. When given,
+                         masked positions are replaced with [MASK] tokens.
+                         Used during pretraining only.
+
+        Returns:
+            patch_out:    (B, N_patches, embed_dim)
+            intermediates: {layer_idx: (B, N_patches, embed_dim)}
+            cls_out:      (B, embed_dim) — global representation
+            global_freq:  (B, 1, embed_dim) — global frequency token
+        """
+        B, C, T = x.shape
+        tokens, global_freq = self.embed(x)
+
+        if mask_matrix is not None:
+            mask_tok = self.mask_token.expand(B, tokens.size(1), -1)
+            tokens = torch.where(mask_matrix.unsqueeze(-1), mask_tok, tokens)
+
+        input_tokens = torch.cat([global_freq, tokens], dim=1)
+        full_out, intermediates, cls_out = self.transformer(input_tokens)
+
+        # Remove global_freq position from patch outputs
+        patch_out = full_out[:, 1:, :]
+        trimmed_intermediates = {k: v[:, 1:, :] for k, v in intermediates.items()}
+
+        return patch_out, trimmed_intermediates, cls_out, global_freq
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Extract CLS token embedding for downstream tasks.
+
+        Args:
+            x: (B, 6, 200) IMU sensor window.
+
+        Returns:
+            (B, embed_dim) global representation.
+        """
+        _, _, cls_out, _ = self.forward(x)
+        return cls_out
+
+    def encode_with_layers(self, x: torch.Tensor) -> Dict[int, torch.Tensor]:
+        """Extract multi-level intermediate representations.
+
+        Args:
+            x: (B, 6, 200) IMU sensor window.
+
+        Returns:
+            {layer_idx: (B, N_patches, embed_dim)} for configured target layers.
+        """
+        _, intermediates, _, _ = self.forward(x)
+        return intermediates
+
+
+# For HF AutoModel compatibility
+IMUMaskedEncoder.config_class = IMUEncoderConfig