diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 560db6a214..78734c8892 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -20,7 +20,6 @@ ) from sdm._inference import inference_mode from sdm._memory import set_pinned_max_round_threshold -from sdm._warnings import warn_once from sdm.cache import Cache from sdm.models.base import Callback, Model from sdm.processing.execution import ( @@ -701,7 +700,8 @@ def _validate_context( f"dimensions (got {tuple(x.size()[:-1])} and " f"{tuple(y.size()[:-1])})" ) - feature_stypes = x.active_stypes + self._validate_feature_stype(x) + if x_context_only is not None: if x_context_only.size()[:-1] != y.size()[:-1]: raise ValueError( @@ -710,18 +710,8 @@ def _validate_context( f"(got {tuple(x_context_only.size()[:-1])} and " f"{tuple(y.size()[:-1])})" ) - feature_stypes = feature_stypes | x_context_only.active_stypes - invalid = feature_stypes - self.supported_feature_stypes - {Stype.id} - if len(invalid) > 0: - stypes = ", ".join(f"{str(stype)!r}" for stype in invalid) - warn_once( - key="model-unsupported-feature-stypes", - message=( - f"{self.__class__.__name__!r} received unsupported " - f"feature stypes {stypes}. Columns with unsupported " - f"feature stypes will not be consumed by the model." - ), - ) + self._validate_feature_stype(x_context_only) + invalid = y.active_stypes - self.supported_target_stypes if len(invalid) > 0: stypes = ", ".join(f"{str(stype)!r}" for stype in invalid) @@ -745,21 +735,8 @@ def _validate_context( f"{self.__class__.__name__!r} does not support related " f"tables" ) - for table_name, table in related_tables.tables.items(): - invalid = table.active_stypes - self.supported_feature_stypes - invalid = invalid - {Stype.id} - if len(invalid) > 0: - stypes = ", ".join(f"{str(stype)!r}" for stype in invalid) - warn_once( - key="model-unsupported-feature-stypes", - message=( - f"{self.__class__.__name__!r} received " - f"unsupported feature stypes {stypes} in related " - f"table {table_name!r}. Columns with unsupported " - f"feature stypes will not be consumed by the " - f"model." - ), - ) + for table in related_tables.tables.values(): + self._validate_feature_stype(table) def _validate_query( self, diff --git a/sdm/models/base/model.py b/sdm/models/base/model.py index a266c9960d..1b48d1ecfd 100644 --- a/sdm/models/base/model.py +++ b/sdm/models/base/model.py @@ -8,7 +8,8 @@ import torch from torch import Tensor -from sdm import Recipe, Stype +from sdm import Recipe, Stype, TableTensor +from sdm._warnings import warn_once from sdm.cache import Cache from sdm.tensor.mixin import DeviceMixin @@ -74,6 +75,21 @@ def __repr__(self) -> str: device_repr = f"device={device}" if device.type != "cpu" else "" return f"{self.__class__.__name__}({device_repr})" + # Helpers ################################################################# + + def _validate_feature_stype(self, x: TableTensor) -> None: + invalid = x.active_stypes - self.supported_feature_stypes - {Stype.id} + if len(invalid) > 0: + stypes = ", ".join(f"{str(stype)!r}" for stype in invalid) + warn_once( + key="model-unsupported-feature-stypes", + message=( + f"{self.__class__.__name__!r} received unsupported " + f"feature stypes {stypes}. Columns with unsupported " + f"feature stypes will not be consumed by the model." + ), + ) + class _Prefetcher(Generic[T]): def __init__( diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 46a3534a7b..15b445e7ca 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -5,6 +5,7 @@ from __future__ import annotations +import copy from collections.abc import Sequence from typing import ClassVar @@ -12,9 +13,11 @@ from torch import Tensor from sdm import Recipe, Stype, TableTensor +from sdm._inference import inference_mode from sdm.models import Callback, Model from sdm.models._huggingface import download_checkpoint from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe +from sdm.processing.execution import MemberQuery, RecipeExecution class KumoAnomaly(Model): @@ -83,14 +86,14 @@ def default_recipe(cls) -> Recipe: def forward( self, - x: Tensor | TableTensor, + x: Tensor | TableTensor, # [ ..., R, D] *, recipe: Recipe | None = None, num_samples: int = 30, num_steps: int = 20, callbacks: Sequence[Callback] | None = None, generator: torch.Generator | None = None, - ) -> TableTensor: + ) -> TableTensor: # [..., R, 1] r"""The forward pass. Args: @@ -111,7 +114,57 @@ def forward( absolute reconstruction error across preprocessed features at each time step. """ - raise NotImplementedError + callbacks = () if callbacks is None else callbacks + requires_grad = self.training + requires_grad |= any(callback.requires_grad for callback in callbacks) + + recipe_execution = RecipeExecution( + self.default_recipe() if recipe is None else copy.deepcopy(recipe) + ) + with ( + torch.amp.autocast(x.device.type, enabled=False), + inference_mode("no_grad" if requires_grad else "inference"), + ): + query = recipe_execution.fit_transform( + x=x, + y=None, + related_tables=None, + num_members=1, + generator=generator, + )[0] + + for callback in callbacks: + query = MemberQuery( + *callback.on_query_preprocessing_end( + model=self, + x=query.x, + related_tables=query.related_tables, + ), + ) + + self._validate_feature_stype(query.x) + + with inference_mode("grad" if requires_grad else "inference"): + out = self.model( + x=query.x.numerical, + num_samples=num_samples, + num_steps=num_steps, + generator=generator, + ) + + out = TableTensor( + columns={Stype.numerical: ("mae",)}, + numerical=out.to(query.x.numerical), + ) + + for callback in callbacks: + out = callback.on_model_forward_end(self, out) + + with ( + torch.amp.autocast(x.device.type, enabled=False), + inference_mode("grad" if requires_grad else "inference"), + ): + return recipe_execution.transform_output([out])[0] class _KumoAnomaly(torch.nn.Module): @@ -121,3 +174,12 @@ def __init__( dtype: torch.dtype | None = None, ) -> None: super().__init__() + + def forward( + self, + x: Tensor, # [..., R, C] + num_samples: int, + num_steps: int, + generator: torch.Generator | None = None, + ) -> Tensor: # [..., R, 1] + return x.mean(dim=-1, keepdim=True) diff --git a/sdm/processing/execution.py b/sdm/processing/execution.py index 16c61f7997..68a4b0e08a 100644 --- a/sdm/processing/execution.py +++ b/sdm/processing/execution.py @@ -5,7 +5,7 @@ import copy from collections.abc import Mapping, Sequence -from typing import NamedTuple, cast +from typing import NamedTuple, cast, overload import torch from torch import Tensor @@ -53,6 +53,7 @@ def num_members(self) -> int: assert self._y_locations is not None return len(self._y_locations) + @overload def fit_transform( self, x: Tensor | TableTensor, @@ -62,60 +63,89 @@ def fit_transform( num_members: int | None = None, context_only_columns: Sequence[str] | Sequence[int] | None = None, generator: torch.Generator | None = None, - ) -> tuple[MemberContext, ...]: + ) -> tuple[MemberContext, ...]: ... + + @overload + def fit_transform( + self, + x: Tensor | TableTensor, + y: None, + related_tables: RelatedTables | None, + *, + num_members: int | None = None, + context_only_columns: None = None, + generator: torch.Generator | None = None, + ) -> tuple[MemberQuery, ...]: ... + + def fit_transform( + self, + x: Tensor | TableTensor, + y: Tensor | TableTensor | None, + related_tables: RelatedTables | None, + *, + num_members: int | None = None, + context_only_columns: Sequence[str] | Sequence[int] | None = None, + generator: torch.Generator | None = None, + ) -> tuple[MemberContext | MemberQuery, ...]: """Fit and transform context data.""" + self._num_estimators = num_members + # Transform target first to be able to resolve task type. Inverse # target transforms require distinct member assignment: - y: EnsembleTable = _to_ensemble_table(y, num_members, expand=True) - - series_processors = tuple( - module - for processor in (self.recipe.features, self.recipe.target) - for module in _iter_modules(processor) - if isinstance(module, sp.SeriesEnsembleProcessor) - ) - for module in series_processors: - module._set_buffer( - name="_offset", - value=torch.zeros( - size=(len(y),), - dtype=torch.long, - device=y[0].device, - ), + series_processors: tuple[sp.SeriesEnsembleProcessor, ...] = () + if y is not None: + y: EnsembleTable = _to_ensemble_table(y, num_members, expand=True) + + series_processors = tuple( + module + for processor in (self.recipe.features, self.recipe.target) + for module in _iter_modules(processor) + if isinstance(module, sp.SeriesEnsembleProcessor) ) + for module in series_processors: + module._set_buffer( + name="_offset", + value=torch.zeros( + size=(len(y),), + dtype=torch.long, + device=y[0].device, + ), + ) - y = self.recipe.target.fit_transform_ensemble(y, generator=generator) - self._num_estimators = num_members - self._y_locations = y._locations + y = self.recipe.target.fit_transform_ensemble( + y, generator=generator + ) + self._y_locations = y._locations - task_dispatchers = tuple( - module - for processor in (self.recipe.features, self.recipe.output) - for module in _iter_modules(processor) - if isinstance(module, sp.TaskDispatch) - ) - if len(task_dispatchers) > 0: - if any(group.size(-1) != 1 for group in y._iter_groups()): - raise ValueError( - "Expected the transformed target to contain exactly one " - "column" - ) + task_dispatchers = tuple( + module + for processor in (self.recipe.features, self.recipe.output) + for module in _iter_modules(processor) + if isinstance(module, sp.TaskDispatch) + ) + if len(task_dispatchers) > 0: + if any(group.size(-1) != 1 for group in y._iter_groups()): + raise ValueError( + "Expected the transformed target to contain exactly " + "one column" + ) - if all( - group.numerical.size(-1) == 1 for group in y._iter_groups() - ): - task = "regression" - elif all( - group.categorical.size(-1) == 1 for group in y._iter_groups() - ): - task = "classification" - else: - raise ValueError( - "'Recipe.target' must resolve to a single task type" - ) + if all( + group.numerical.size(-1) == 1 for group in y._iter_groups() + ): + task = "regression" + elif all( + group.categorical.size(-1) == 1 + for group in y._iter_groups() + ): + task = "classification" + else: + raise ValueError( + "'Recipe.target' must resolve to a single task type" + ) - for task_dispatcher in task_dispatchers: - task_dispatcher._task = task + for task_dispatcher in task_dispatchers: + task_dispatcher._task = task self._related_processors = None self._related_locations = None @@ -135,11 +165,6 @@ def fit_transform( ensemble_table, generator=generator, ) - if len(related_ensembles[name]) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) for module in _iter_modules(processor): if isinstance(module, sp.SeriesEnsembleProcessor): module._set_buffer( @@ -160,6 +185,7 @@ def fit_transform( x_context_only: EnsembleTable | None = None if context_only_columns: + assert y is not None if not isinstance(x, TableTensor): x = TableTensor.from_tensor(cast(Tensor, x)) context_only_columns = [str(col) for col in context_only_columns] @@ -170,11 +196,6 @@ def fit_transform( ), generator=generator, ) - if len(x_context_only) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) x = x.drop_columns(context_only_columns) inputs = _to_ensemble_table(x, num_members) @@ -182,10 +203,6 @@ def fit_transform( inputs, generator=generator, ) - if len(x) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of ensemble members" - ) for module in series_processors: module._set_buffer( @@ -197,8 +214,8 @@ def fit_transform( ), ) - members: list[MemberContext] = [] - for member_id in range(self.num_members): + members: list[MemberContext | MemberQuery] = [] + for member_id in range(len(x)): related_tables_i: RelatedTables[TableTensor] | None = None if related_tables is not None: related_tables_i = RelatedTables( @@ -209,8 +226,8 @@ def fit_transform( relationships=related_tables.relationships, task_links=related_tables.task_links, ) - members.append( - MemberContext( + if y is not None: + member = MemberContext( x=x[member_id], y=y[member_id], x_context_only=x_context_only[member_id] @@ -219,7 +236,12 @@ def fit_transform( related_tables=related_tables_i, input_stypes=inputs[member_id].stypes, ) - ) + else: + member = MemberQuery( + x=x[member_id], + related_tables=related_tables_i, + ) + members.append(member) return tuple(members) @@ -231,10 +253,6 @@ def transform( """Transform query data.""" x: EnsembleTable = _to_ensemble_table(x, self._num_estimators) x = self.recipe.features.transform_ensemble(x) - if len(x) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of ensemble members" - ) related_ensembles: Mapping[str, EnsembleTable] = {} if related_tables is not None: @@ -252,14 +270,9 @@ def transform( related_ensembles[name] = processor.transform_ensemble( ensemble_table ) - if len(related_ensembles[name]) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) members: list[MemberQuery] = [] - for member_id in range(self.num_members): + for member_id in range(len(x)): related_tables_i: RelatedTables[TableTensor] | None = None if related_tables is not None: related_tables_i = RelatedTables( diff --git a/sdm/processing/numerical/min_max_scale.py b/sdm/processing/numerical/min_max_scale.py index 4b68a78d4e..1915093459 100644 --- a/sdm/processing/numerical/min_max_scale.py +++ b/sdm/processing/numerical/min_max_scale.py @@ -10,7 +10,7 @@ class MinMaxScale(Processor, InvertibleMixin): - """Scale each numerical column by its fitted minimum and range. + """Scale each numerical column by its fitted min-max range. NaN and infinite values are ignored when fitting statistics and preserved during the transform. diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index c1e37e9ee7..fee8df3510 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -63,3 +63,9 @@ def _transform(self, table: TableTensor) -> TableTensor: TableTensor, torch.cat([table.drop_stypes(Stype.numerical), out], dim=-1), ) + + def __repr__(self, *, indent: int = 0) -> str: + return ( + f"{' ' * indent}{self.__class__.__name__}(" + f"num_components={self.num_components})" + ) diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py new file mode 100644 index 0000000000..996237d59b --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm import TableTensor +from sdm.models import KumoAnomaly +from sdm.testing import withCUDA + + +@withCUDA +def test_forward(device: torch.device) -> None: + model = KumoAnomaly(pretrained=False, device=device) + x = TableTensor.from_tensor(torch.randn(32, 5, device=device)) + + out = model(x) + assert out.size() == (32, 1) + assert out.device == device + assert out.numerical.isfinite().all() + assert "mae" in out.numerical_columns