diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py index f9e4931c6..dc69ce792 100644 --- a/sdm/models/kumo/timeseries/anomaly/denoiser.py +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -29,6 +29,7 @@ def __init__( ) -> None: super().__init__() factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.num_steps = num_steps self.diffusion_embedding = DiffusionEmbedding( num_steps=num_steps, diff --git a/sdm/models/kumo/timeseries/anomaly/embedding.py b/sdm/models/kumo/timeseries/anomaly/embedding.py index 5c2feae86..f9f154bc9 100644 --- a/sdm/models/kumo/timeseries/anomaly/embedding.py +++ b/sdm/models/kumo/timeseries/anomaly/embedding.py @@ -87,7 +87,7 @@ class SideEmbedding(torch.nn.Module): before conversion to the parameter dtype. Args: - num_features: Number of features in checkpoint order. + in_channels: Number of features in checkpoint order. time_channels: Even sinusoidal time embedding width, at least two. feature_channels: Learned feature embedding width. is_unconditional: Append the conditioning mask as the last channel. @@ -97,7 +97,7 @@ class SideEmbedding(torch.nn.Module): def __init__( self, - num_features: int, + in_channels: int, time_channels: int = 128, feature_channels: int = 16, is_unconditional: bool = False, @@ -109,7 +109,7 @@ def __init__( self.time_channels = time_channels self.is_unconditional = is_unconditional self.embed_layer = torch.nn.Embedding( - num_embeddings=num_features, + num_embeddings=in_channels, embedding_dim=feature_channels, **factory_kwargs, ) @@ -129,7 +129,7 @@ def forward(self, time: Tensor, mask: Tensor) -> Tensor: time: Model time coordinates with shape ``[B, L]``. Coordinates are used as supplied, with no normalization or resampling. mask: Conditioning mask with shape ``[B, K, L]`` in checkpoint - feature order, where ``K`` equals ``num_features``. One marks + feature order, where ``K`` equals ``in_channels``. One marks an observed value used for conditioning; zero marks a target. Both inputs must be on the module's device. diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 15b445e7c..b4d7929ea 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -7,16 +7,22 @@ import copy from collections.abc import Sequence -from typing import ClassVar +from functools import partial +from typing import Any, ClassVar, cast import torch +import torch.nn.functional as F 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.denoiser import DiffusionDenoiser +from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule from sdm.processing.execution import MemberQuery, RecipeExecution @@ -154,12 +160,14 @@ def forward( out = TableTensor( columns={Stype.numerical: ("mae",)}, - numerical=out.to(query.x.numerical), + numerical=out.unsqueeze(-1), ) for callback in callbacks: out = callback.on_model_forward_end(self, out) + out = cast(TableTensor, out.to(query.x.dtype)) + with ( torch.amp.autocast(x.device.type, enabled=False), inference_mode("grad" if requires_grad else "inference"), @@ -170,16 +178,137 @@ def forward( class _KumoAnomaly(torch.nn.Module): def __init__( self, + in_channels: int = 40, + time_channels: int = 256, + feature_channels: int = 64, + channels: int = 128, + num_layers: int = 6, + num_heads: int = 8, + embedding_channels: int = 256, + num_steps: int = 500, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.in_channels = in_channels + + self.side_embedding = SideEmbedding( + in_channels=in_channels, + time_channels=time_channels, + feature_channels=feature_channels, + **factory_kwargs, + ) + self.denoiser = DiffusionDenoiser( + channels=channels, + side_channels=time_channels + feature_channels, + num_steps=num_steps, + num_layers=num_layers, + num_heads=num_heads, + embedding_channels=embedding_channels, + **factory_kwargs, + ) def forward( self, - x: Tensor, # [..., R, C] + x: Tensor, # [..., R, D] num_samples: int, num_steps: int, + window_size: int = 100, + num_blocks: int = 4, + batch_size: int = 32, generator: torch.Generator | None = None, - ) -> Tensor: # [..., R, 1] - return x.mean(dim=-1, keepdim=True) + ) -> Tensor: # [..., R] + + if x.size(-1) > self.in_channels: + raise ValueError( + f"'KumoAnomaly' supports at most {self.in_channels} columns " + f"after pre-processing (got {x.size(-1)})" + ) + + *B, R, D = x.size() + x = F.pad(x.reshape(-1, R, D), (0, self.in_channels - D)) + + # Pad last window by repeating last row: + x = torch.cat([x, x[:, -1:].expand(-1, -R % window_size, -1)], dim=-2) + window = x.reshape(-1, window_size, self.in_channels) + + schedule = DiffusionSchedule( + num_steps=self.denoiser.num_steps, + beta_start=0.0001, + beta_end=0.01, + schedule="linear", + device="cpu", + ) + sampler = DPMSolver(schedule.beta.float(), num_steps).to(x.device) + + out = torch.cat( + [ + self.reconstruct( + x=_x, + sampler=sampler, + num_samples=num_samples, + num_blocks=num_blocks, + generator=generator, + ) + for _x in window.split(batch_size) + ], + dim=0, + ) + + out = out.view(x.size())[:, :R, :D] + out = (out - x[:, :R, :D]).abs().mean(dim=-1) + return out.view(*B, R) + + def reconstruct( + self, + x: Tensor, # [B, R, D] + sampler: DPMSolver, + num_samples: int, + num_blocks: int, + generator: torch.Generator | None, + ) -> Tensor: # [B, R, D] + + B, R, _ = x.size() + x = x.transpose(1, 2) # [B, D, R] + mean = x.mean(dim=(1, 2), keepdim=True) + std = x.std(dim=(1, 2), keepdim=True) + 1e-5 + x = ((x - mean) / std).clamp(-10, 10) + + time = torch.arange(R, device=x.device).expand(B, -1) + blocks = time // (R // num_blocks) + + out = x.new_zeros((num_samples, *x.size())) + for strategy in range(2): + mask = ((blocks % 2) == strategy).to(x.dtype) + mask = mask.unsqueeze(1).expand_as(x) # [B, D, R] + predict_noise = partial( + self.predict_noise, + cond=mask * x, + target_mask=1 - mask, + side_info=self.side_embedding(time, mask), + strategy=torch.tensor([strategy], device=x.device), + ) + for i in range(num_samples): + noise = torch.empty_like(x).normal_(generator=generator) + sample = sampler(noise, predict_noise) * std + mean + out[i] += (1 - mask) * sample + + return out.median(dim=0).values.transpose(1, 2) + + def predict_noise( + self, + x: Tensor, # [B, D, R] + step: Tensor, # [1] + *, + cond: Tensor, # [B, D, R] + target_mask: Tensor, # [B, D, R] + side_info: Tensor, # [B, S, D, R] + strategy: Tensor, # [1] + ) -> Tensor: # [B, D, R] + return self.denoiser( + x=torch.stack([cond, target_mask * x], dim=1), + side_info=side_info, + step=step, + strategy=strategy, + ) diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py index 996237d59..849dd2874 100644 --- a/test/models/kumo/timeseries/anomaly/test_model.py +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -3,7 +3,6 @@ import torch -from sdm import TableTensor from sdm.models import KumoAnomaly from sdm.testing import withCUDA @@ -11,10 +10,11 @@ @withCUDA def test_forward(device: torch.device) -> None: model = KumoAnomaly(pretrained=False, device=device) - x = TableTensor.from_tensor(torch.randn(32, 5, device=device)) + x = torch.randn(2, 130, 3, device=device) - out = model(x) - assert out.size() == (32, 1) + out = model(x, num_samples=1, num_steps=1) + assert out.size() == (2, 130, 1) + assert out.columns["numerical"] == ("mae",) assert out.device == device assert out.numerical.isfinite().all() - assert "mae" in out.numerical_columns + assert (out.numerical >= 0).all() diff --git a/test/models/kumo/timeseries/anomaly/test_side_embedding.py b/test/models/kumo/timeseries/anomaly/test_side_embedding.py index 0c12b4cdf..76961164e 100644 --- a/test/models/kumo/timeseries/anomaly/test_side_embedding.py +++ b/test/models/kumo/timeseries/anomaly/test_side_embedding.py @@ -35,9 +35,9 @@ def test_side_embedding( and not torch.cuda.is_bf16_supported() ): pytest.skip("CUDA bfloat16 not supported") - batch, features, length = shape + B, D, R = shape module = SideEmbedding( - num_features=features, + in_channels=D, time_channels=time_channels, feature_channels=4, is_unconditional=is_unconditional, @@ -46,10 +46,10 @@ def test_side_embedding( ) if dtype == torch.float64: module = module.to(dtype=dtype) - time = torch.arange(batch * length, device=device).reshape(batch, length) + time = torch.arange(B * R, device=device).reshape(B, R) time = time.float().square() * 0.25 - 2.5 time[:, -1] = 100000 - mask = torch.arange(batch * features * length, device=device) + mask = torch.arange(B * D * R, device=device) mask = mask.reshape(shape).remainder(2).bool() actual = module(time, mask) @@ -60,7 +60,7 @@ def test_side_embedding( / time_channels, ) phase = time.unsqueeze(-1) * frequencies - time_embed = torch.empty(batch, length, time_channels, device=device) + time_embed = torch.empty(B, R, time_channels, device=device) time_embed[..., 0::2] = phase.sin() time_embed[..., 1::2] = phase.cos() time_embed = F.layer_norm( @@ -71,7 +71,7 @@ def test_side_embedding( eps=module.time_embed_norm.eps, ) feature_embed = F.embedding( - input=torch.arange(features, device=device), + input=torch.arange(D, device=device), weight=module.embed_layer.weight, ) feature_embed = F.layer_norm( @@ -83,19 +83,14 @@ def test_side_embedding( ) expected = torch.cat( [ - time_embed[:, :, None].expand(-1, -1, features, -1), - feature_embed[None, None].expand(batch, length, -1, -1), + time_embed[:, :, None].expand(-1, -1, D, -1), + feature_embed[None, None].expand(B, R, -1, -1), ], dim=-1, ).permute(0, 3, 2, 1) if is_unconditional: expected = torch.cat([expected, mask[:, None].to(dtype)], dim=1) - assert actual.shape == ( - batch, - time_channels + 4 + int(is_unconditional), - features, - length, - ) + assert actual.shape == (B, time_channels + 4 + int(is_unconditional), D, R) assert actual.dtype == dtype assert actual.device == device torch.testing.assert_close(actual, expected)