| |
| |
|
|
| """IMU-only COMODO inference, preserving the original checkpoint parameter names.""" |
|
|
| from dataclasses import dataclass |
| from typing import Optional |
|
|
| import torch |
| from mantis.architecture import MantisV2 |
| from torch import nn |
| from torch.nn import functional as F |
| from transformers import PreTrainedModel |
| from transformers.utils import ModelOutput |
|
|
| from .configuration_comodo import ComodoConfig |
|
|
|
|
| @dataclass |
| class ComodoOutput(ModelOutput): |
| embeddings: Optional[torch.Tensor] = None |
| projection: Optional[torch.Tensor] = None |
|
|
|
|
| class ComodoBundle(PreTrainedModel): |
| """Resolve the root bundle before loading a selected, fully trained model.""" |
|
|
| config_class = ComodoConfig |
| main_input_name = "imu_values" |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *args, config=None, **kwargs): |
| hub_keys = ("revision", "token", "cache_dir", "local_files_only", "force_download") |
| hub_kwargs = {k: kwargs[k] for k in hub_keys if k in kwargs} |
| config = config or cls.config_class.from_pretrained(pretrained_model_name_or_path, **hub_kwargs) |
| if config._commit_hash is not None: |
| hub_kwargs["revision"] = config._commit_hash |
| load_kwargs = {k: v for k, v in kwargs.items() |
| if k not in ("variant", "_from_auto", "_from_pipeline", "adapter_kwargs")} |
| load_kwargs.update(hub_kwargs) |
| if "torch_dtype" not in load_kwargs and config.torch_dtype is not None: |
| load_kwargs["torch_dtype"] = config.torch_dtype |
| if kwargs.get("_from_pipeline", getattr(config, "_release_from_pipeline", None)) == "comodo": |
| |
| |
| selector = cls(config) |
| selector.name_or_path = str(pretrained_model_name_or_path) |
| selector._release_load_kwargs = load_kwargs |
| return selector |
| variant = kwargs.pop("variant", None) or config.variant |
| if variant not in ("core50", "fine150"): |
| raise ValueError("variant must be 'core50' or 'fine150'") |
| return ComodoModel.from_pretrained(pretrained_model_name_or_path, *args, |
| subfolder=variant, **load_kwargs) |
|
|
| def forward(self, *args, **kwargs): |
| raise RuntimeError("Select a variant through pipeline('comodo', ..., variant='core50')") |
|
|
|
|
| class ComodoModel(PreTrainedModel): |
| config_class = ComodoConfig |
| main_input_name = "imu_values" |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| |
| |
| self.imu_student = MantisV2(**config.mantis_config, device="cpu") |
| self.mlp = nn.Sequential( |
| nn.Linear(config.feature_dim, config.projection_hidden_dim), |
| nn.BatchNorm1d(config.projection_hidden_dim), |
| nn.GELU(), |
| nn.Linear(config.projection_hidden_dim, config.projection_dim), |
| ) |
| self.post_init() |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): |
| subfolder = kwargs.get("subfolder", "") |
| hub_kwargs = {k: kwargs[k] for k in |
| ("revision", "token", "cache_dir", "local_files_only") if k in kwargs} |
| model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) |
| target = model[0] if isinstance(model, tuple) else model |
| target._release_subfolder = subfolder |
| target._release_hub_kwargs = hub_kwargs |
| return model |
|
|
| def forward(self, imu_values, normalize_embeddings=True, return_projection=False, |
| return_dict=True): |
| if imu_values.ndim != 3 or imu_values.shape[1] != self.config.n_channels: |
| raise ValueError(f"Expected [B, {self.config.n_channels}, T] IMU values") |
| if imu_values.shape[-1] != self.config.sequence_length: |
| imu_values = F.interpolate(imu_values, size=self.config.sequence_length, |
| mode="linear", align_corners=False) |
| features = torch.cat([ |
| self.imu_student(imu_values[:, ch:ch + 1, :]) |
| for ch in range(self.config.n_channels) |
| ], dim=1) |
| projection = self.mlp(features) if return_projection else None |
| embeddings = F.normalize(features, p=2, dim=1) if normalize_embeddings else features |
| if not return_dict: |
| return (embeddings, projection) if return_projection else (embeddings,) |
| return ComodoOutput(embeddings=embeddings, projection=projection) |
|
|
|
|
| ComodoModel.register_for_auto_class("AutoModel") |
|
|
