diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py new file mode 100644 index 000000000..f9e4931c6 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: D101, D102 + +import math +from typing import Any + +import torch +from torch import Tensor +from torch.nn import Embedding, LayerNorm, Linear, ModuleList + +from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock +from sdm.models.kumo.timeseries.anomaly.embedding import DiffusionEmbedding + + +class DiffusionDenoiser(torch.nn.Module): + def __init__( + self, + channels: int, + side_channels: int, + num_steps: int, + num_layers: int, + num_heads: int, + embedding_channels: int, + input_channels: int = 2, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + + self.diffusion_embedding = DiffusionEmbedding( + num_steps=num_steps, + channels=embedding_channels, + **factory_kwargs, + ) + self.strategy_embedding = Embedding( + num_embeddings=2, + embedding_dim=embedding_channels, + **factory_kwargs, + ) + self.input_projection = Linear( + input_channels, + channels, + **factory_kwargs, + ) + self.residual_layers = ModuleList( + [ + ResidualBlock( + channels=channels, + side_channels=side_channels, + embedding_channels=embedding_channels, + num_heads=num_heads, + **factory_kwargs, + ) + for _ in range(num_layers) + ] + ) + self.norm_after_skip = LayerNorm(channels, **factory_kwargs) + self.output_projection1 = Linear(channels, channels, **factory_kwargs) + self.norm_after_proj1 = LayerNorm(channels, **factory_kwargs) + self.output_projection2 = Linear(channels, 1, **factory_kwargs) + + for lin in (self.input_projection, self.output_projection1): + torch.nn.init.kaiming_normal_( + tensor=lin.weight, + mode="fan_out", + nonlinearity="relu", + ) + with torch.no_grad(): + lin.weight.mul_(0.5) + torch.nn.init.xavier_uniform_( + tensor=self.output_projection2.weight, + gain=0.01, + ) + + def forward( + self, + x: Tensor, # [B, I, K, L] + side_info: Tensor, # [B, S, K, L] + step: Tensor, # [B] or [1] + strategy: Tensor, # [B] or [1] + ) -> Tensor: # [B, K, L] + x = self.input_projection(x.permute(0, 2, 3, 1)).relu() # [B, K, L, C] + x = x.permute(0, 3, 1, 2) # [B, C, K, L] + + diffusion_embedding = self.diffusion_embedding(step) + strategy_embedding = self.strategy_embedding(strategy) + + skips = [] + for layer in self.residual_layers: + x, skip = layer( + x=x, + side_info=side_info, + diffusion_embedding=diffusion_embedding, + strategy_embedding=strategy_embedding, + ) + skips.append(skip) + + x = torch.stack(skips).sum(dim=0) # [B, C, K, L] + x = x / math.sqrt(len(skips)) + x = self.norm_after_skip(x.permute(0, 2, 3, 1)) # [B, K, L, C] + x = self.output_projection1(x) + x = self.norm_after_proj1(x).relu() + return self.output_projection2(x).squeeze(-1) # [B, K, L] diff --git a/test/models/kumo/timeseries/anomaly/test_denoiser.py b/test/models/kumo/timeseries/anomaly/test_denoiser.py new file mode 100644 index 000000000..bed3cadf7 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_denoiser.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize("input_channels", [1, 2]) +@pytest.mark.parametrize(("features", "length"), [(3, 5), (1, 1)]) +def test_diffusion_denoiser( + input_channels: int, + features: int, + length: int, + device: torch.device, +) -> None: + model = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=2, + num_heads=2, + embedding_channels=12, + input_channels=input_channels, + device=device, + ) + + x = torch.randn(2, input_channels, features, length, device=device) + side_info = torch.randn(2, 6, features, length, device=device) + step = torch.tensor([0, 9], device=device) + strategy = torch.tensor([1], device=device) + + with torch.no_grad(): + out = model(x, side_info, step, strategy) + + assert out.size() == (2, features, length) + assert out.dtype == x.dtype + assert out.device == device