From 8164ccede210d02d8768b27221ee173048a6db3f Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Fri, 9 Oct 2026 09:51:37 -0700 Subject: [PATCH 1/4] Add Kumo-Anomaly diffusion denoiser with focused tests Signed-off-by: Aditi Gautam --- .../kumo/timeseries/anomaly/denoiser.py | 148 +++++++++++++++ .../kumo/timeseries/anomaly/test_denoiser.py | 171 ++++++++++++++++++ 2 files changed, 319 insertions(+) create mode 100644 sdm/models/kumo/timeseries/anomaly/denoiser.py create mode 100644 test/models/kumo/timeseries/anomaly/test_denoiser.py diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py new file mode 100644 index 0000000000..eaf22276b7 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import torch +from torch import Tensor +from torch.nn import Conv1d, Embedding, LayerNorm, ModuleList + +from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock +from sdm.models.kumo.timeseries.anomaly.embedding import DiffusionEmbedding + + +class DiffusionDenoiser(torch.nn.Module): + """Predict noise for one Kumo-Anomaly diffusion step. + + Callers supply normalized noisy/observed inputs and precomputed side + information. This component does not sample or compute anomaly scores. + + Args: + channels: Width of the residual blocks. + side_channels: Width of the precomputed conditioning information. + num_steps: Number of discrete diffusion steps. + num_layers: Number of residual blocks, at least one. + num_heads: Attention heads for each axis. + embedding_channels: Even diffusion and strategy embedding width, + at least four. + input_channels: Input width (two for conditional, one unconditional). + hidden_channels: Transformer feedforward width. + dropout: Transformer attention and feedforward dropout probability. + device: The device. + dtype: The parameter and output dtype. + """ + + def __init__( + self, + channels: int, + side_channels: int, + num_steps: int, + num_layers: int, + num_heads: int, + embedding_channels: int = 128, + input_channels: int = 2, + hidden_channels: int = 64, + dropout: float = 0.1, + 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 = Conv1d( + in_channels=input_channels, + out_channels=channels, + kernel_size=1, + **factory_kwargs, + ) + self.output_projection1 = Conv1d( + in_channels=channels, + out_channels=channels, + kernel_size=1, + **factory_kwargs, + ) + self.output_projection2 = Conv1d( + in_channels=channels, + out_channels=1, + kernel_size=1, + **factory_kwargs, + ) + for projection in (self.input_projection, self.output_projection1): + torch.nn.init.kaiming_normal_( + tensor=projection.weight, mode="fan_out", nonlinearity="relu" + ) + with torch.no_grad(): + projection.weight.mul_(0.5) + torch.nn.init.xavier_uniform_( + tensor=self.output_projection2.weight, gain=0.01 + ) + self.residual_layers = ModuleList( + ResidualBlock( + channels=channels, + side_channels=side_channels, + embedding_channels=embedding_channels, + num_heads=num_heads, + hidden_channels=hidden_channels, + dropout=dropout, + **factory_kwargs, + ) + for _ in range(num_layers) + ) + self.norm_after_skip = LayerNorm(channels, **factory_kwargs) + self.norm_after_proj1 = LayerNorm(channels, **factory_kwargs) + + def forward( + self, + x: Tensor, + side_info: Tensor, + step: Tensor, + strategy: Tensor, + ) -> Tensor: + """Predict noise for normalized time-series windows. + + Args: + x: Input with shape ``[B, I, K, L]``, where ``B`` is batch size, + ``I`` input channels, ``K`` features, and ``L`` time steps. + side_info: Conditioning with shape ``[B, S, K, L]``, where + ``S`` is ``side_channels``. + step: Integer diffusion indices in ``[0, num_steps)`` with + shape ``[B]`` or ``[1]``. + strategy: Integer masking-strategy indices in ``{0, 1}`` with + shape ``[B]`` or ``[1]``. Single-entry indices broadcast + across the batch independently for either embedding. + + Returns: + Predicted noise with shape ``[B, K, L]`` in input feature/time + order. No denormalization or mask restoration is applied. + """ + batch, _, features, length = x.shape + x = self.input_projection(x.flatten(2)).relu() + x = x.reshape(batch, -1, features, length) + 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) / math.sqrt(len(skips)) + x = self.norm_after_skip(x.flatten(2).transpose(1, 2)) # [B, K * L, C] + x = self.output_projection1(x.transpose(1, 2)) + x = self.norm_after_proj1(x.transpose(1, 2)).transpose(1, 2) + return self.output_projection2(x.relu()).reshape( + batch, features, length + ) 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 0000000000..399fc6e7c4 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_denoiser.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +import pytest +import torch +import torch.nn.functional as F +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser +from sdm.testing import withCUDA + + +def reference_denoiser( + model: DiffusionDenoiser, + x: Tensor, + side: Tensor, + step: Tensor, + strategy: Tensor, +) -> Tensor: + """Compose the tested blocks with pointwise linear projections.""" + x = x.permute(0, 2, 3, 1) + projection = model.input_projection + x = F.linear(x, projection.weight.squeeze(-1), projection.bias).relu() + x = x.permute(0, 3, 1, 2) + diffusion = model.diffusion_embedding(step) + strategy_embedding = model.strategy_embedding(strategy) + skips = [] + for block in model.residual_layers: + x, skip = block(x, side, diffusion, strategy_embedding) + skips.append(skip.permute(0, 2, 3, 1)) + x = sum(skips) / math.sqrt(len(skips)) + x = model.norm_after_skip(x) + projection = model.output_projection1 + x = F.linear(x, projection.weight.squeeze(-1), projection.bias) + x = model.norm_after_proj1(x).relu() + projection = model.output_projection2 + x = F.linear(x, projection.weight.squeeze(-1), projection.bias) + return x.squeeze(-1) + + +@withCUDA +@pytest.mark.parametrize("input_channels", [1, 2]) +@pytest.mark.parametrize( + ("features", "length"), [(3, 5), (1, 5), (3, 1), (1, 1)] +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("num_layers", [1, 3]) +@pytest.mark.parametrize(("step_batch", "strategy_batch"), [(1, 2), (2, 1)]) +def test_denoiser( + device: torch.device, + input_channels: int, + features: int, + length: int, + dtype: torch.dtype, + num_layers: int, + step_batch: int, + strategy_batch: int, +) -> None: + """Match pointwise reference outputs and independent window predictions.""" + model = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=num_layers, + num_heads=2, + embedding_channels=12, + input_channels=input_channels, + device=device, + dtype=dtype, + ).eval() + x = torch.randn( + 2, input_channels, features, length, device=device, dtype=dtype + ) + side = torch.randn(2, 6, features, length, device=device, dtype=dtype) + step = torch.tensor([0, 9], device=device)[:step_batch] + strategy = torch.tensor([0, 1], device=device)[:strategy_batch] + with torch.inference_mode(): + actual = model(x, side, step, strategy) + expected = reference_denoiser(model, x, side, step, strategy) + assert actual.shape == (2, features, length) + assert actual.dtype == dtype + assert actual.device == device + assert actual.isfinite().all() + torch.testing.assert_close(actual, expected) + individual = torch.cat( + [ + model( + x=x[i : i + 1], + side_info=side[i : i + 1], + step=step if step_batch == 1 else step[i : i + 1], + strategy=strategy + if strategy_batch == 1 + else strategy[i : i + 1], + ) + for i in range(2) + ] + ) + torch.testing.assert_close(actual, individual) + + +@withCUDA +def test_denoiser_grad_and_state_dict(device: torch.device) -> None: + """Propagate input gradients and restore saved predictions.""" + model = ( + DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=2, + num_heads=2, + embedding_channels=12, + ) + .to(device=device, dtype=torch.float64) + .eval() + ) + x = torch.randn( + 2, 2, 3, 5, device=device, dtype=torch.float64, requires_grad=True + ) + side = torch.randn( + 2, 6, 3, 5, device=device, dtype=torch.float64, requires_grad=True + ) + step = torch.tensor([0, 9], device=device) + strategy = torch.tensor([0, 1], device=device) + out = model(x, side, step, strategy) + out.square().mean().backward() + for value in (x, side, *model.parameters()): + assert value.grad is not None + assert value.grad.isfinite().all() + for value in (x, side): + assert value.grad is not None + assert value.grad.abs().sum() > 0 + restored = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=2, + num_heads=2, + embedding_channels=12, + device=device, + dtype=torch.float64, + ).eval() + restored.load_state_dict(model.state_dict(), strict=True) + torch.testing.assert_close(restored(x, side, step, strategy), out) + + +@withCUDA +def test_denoiser_dropout(device: torch.device) -> None: + """Apply training dropout and return repeatable evaluation predictions.""" + model = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=2, + num_heads=2, + embedding_channels=12, + dropout=1.0, + device=device, + ) + x = torch.randn(2, 2, 3, 5, device=device) + side = torch.randn(2, 6, 3, 5, device=device) + step = torch.tensor([2, 7], device=device) + strategy = torch.tensor([0, 1], device=device) + with torch.no_grad(): + training = model(x, side, step, strategy) + assert training.isfinite().all() + model.eval() + evaluation = model(x, side, step, strategy) + assert not torch.allclose(training, evaluation) + torch.testing.assert_close(model(x, side, step, strategy), evaluation) From 063b2af862927ec7e75181167d6ea70446a43372 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:27:35 +0900 Subject: [PATCH 2/4] Trim diffusion denoiser tests to focused cases Signed-off-by: Ard <88337265+aw471@users.noreply.github.com> --- .../kumo/timeseries/anomaly/test_denoiser.py | 107 +++++------------- 1 file changed, 28 insertions(+), 79 deletions(-) diff --git a/test/models/kumo/timeseries/anomaly/test_denoiser.py b/test/models/kumo/timeseries/anomaly/test_denoiser.py index 399fc6e7c4..f18f310b6c 100644 --- a/test/models/kumo/timeseries/anomaly/test_denoiser.py +++ b/test/models/kumo/timeseries/anomaly/test_denoiser.py @@ -41,24 +41,30 @@ def reference_denoiser( @withCUDA -@pytest.mark.parametrize("input_channels", [1, 2]) @pytest.mark.parametrize( - ("features", "length"), [(3, 5), (1, 5), (3, 1), (1, 1)] + ( + "input_channels", + "num_layers", + "features", + "length", + "step_batch", + "strategy_batch", + ), + [ + (2, 1, 3, 5, 1, 2), + (1, 3, 3, 5, 2, 1), + (2, 2, 1, 1, 2, 2), + ], ) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("num_layers", [1, 3]) -@pytest.mark.parametrize(("step_batch", "strategy_batch"), [(1, 2), (2, 1)]) def test_denoiser( device: torch.device, input_channels: int, features: int, length: int, - dtype: torch.dtype, num_layers: int, step_batch: int, strategy_batch: int, ) -> None: - """Match pointwise reference outputs and independent window predictions.""" model = DiffusionDenoiser( channels=8, side_channels=6, @@ -68,53 +74,33 @@ def test_denoiser( embedding_channels=12, input_channels=input_channels, device=device, - dtype=dtype, ).eval() - x = torch.randn( - 2, input_channels, features, length, device=device, dtype=dtype - ) - side = torch.randn(2, 6, features, length, device=device, dtype=dtype) + x = torch.randn(2, input_channels, features, length, device=device) + side = torch.randn(2, 6, features, length, device=device) step = torch.tensor([0, 9], device=device)[:step_batch] strategy = torch.tensor([0, 1], device=device)[:strategy_batch] with torch.inference_mode(): actual = model(x, side, step, strategy) expected = reference_denoiser(model, x, side, step, strategy) assert actual.shape == (2, features, length) - assert actual.dtype == dtype + assert actual.dtype == x.dtype assert actual.device == device assert actual.isfinite().all() torch.testing.assert_close(actual, expected) - individual = torch.cat( - [ - model( - x=x[i : i + 1], - side_info=side[i : i + 1], - step=step if step_batch == 1 else step[i : i + 1], - strategy=strategy - if strategy_batch == 1 - else strategy[i : i + 1], - ) - for i in range(2) - ] - ) - torch.testing.assert_close(actual, individual) @withCUDA -def test_denoiser_grad_and_state_dict(device: torch.device) -> None: - """Propagate input gradients and restore saved predictions.""" - model = ( - DiffusionDenoiser( - channels=8, - side_channels=6, - num_steps=10, - num_layers=2, - num_heads=2, - embedding_channels=12, - ) - .to(device=device, dtype=torch.float64) - .eval() - ) +def test_denoiser_grad(device: torch.device) -> None: + model = DiffusionDenoiser( + channels=8, + side_channels=6, + num_steps=10, + num_layers=2, + num_heads=2, + embedding_channels=12, + device=device, + dtype=torch.float64, + ).eval() x = torch.randn( 2, 2, 3, 5, device=device, dtype=torch.float64, requires_grad=True ) @@ -124,6 +110,7 @@ def test_denoiser_grad_and_state_dict(device: torch.device) -> None: step = torch.tensor([0, 9], device=device) strategy = torch.tensor([0, 1], device=device) out = model(x, side, step, strategy) + assert out.dtype == torch.float64 out.square().mean().backward() for value in (x, side, *model.parameters()): assert value.grad is not None @@ -131,41 +118,3 @@ def test_denoiser_grad_and_state_dict(device: torch.device) -> None: for value in (x, side): assert value.grad is not None assert value.grad.abs().sum() > 0 - restored = DiffusionDenoiser( - channels=8, - side_channels=6, - num_steps=10, - num_layers=2, - num_heads=2, - embedding_channels=12, - device=device, - dtype=torch.float64, - ).eval() - restored.load_state_dict(model.state_dict(), strict=True) - torch.testing.assert_close(restored(x, side, step, strategy), out) - - -@withCUDA -def test_denoiser_dropout(device: torch.device) -> None: - """Apply training dropout and return repeatable evaluation predictions.""" - model = DiffusionDenoiser( - channels=8, - side_channels=6, - num_steps=10, - num_layers=2, - num_heads=2, - embedding_channels=12, - dropout=1.0, - device=device, - ) - x = torch.randn(2, 2, 3, 5, device=device) - side = torch.randn(2, 6, 3, 5, device=device) - step = torch.tensor([2, 7], device=device) - strategy = torch.tensor([0, 1], device=device) - with torch.no_grad(): - training = model(x, side, step, strategy) - assert training.isfinite().all() - model.eval() - evaluation = model(x, side, step, strategy) - assert not torch.allclose(training, evaluation) - torch.testing.assert_close(model(x, side, step, strategy), evaluation) From dd8b5739d0327383ce521af3ce1e6fa821344d77 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:14:46 -0700 Subject: [PATCH 3/4] Simplify `DiffusionDenoiser` with channels-last linear projections --- .../kumo/timeseries/anomaly/denoiser.py | 134 ++++++------------ .../kumo/timeseries/anomaly/test_denoiser.py | 109 ++------------ 2 files changed, 61 insertions(+), 182 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py index eaf22276b7..371faa7776 100644 --- a/sdm/models/kumo/timeseries/anomaly/denoiser.py +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -1,38 +1,20 @@ # 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 Conv1d, Embedding, LayerNorm, ModuleList +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): - """Predict noise for one Kumo-Anomaly diffusion step. - - Callers supply normalized noisy/observed inputs and precomputed side - information. This component does not sample or compute anomaly scores. - - Args: - channels: Width of the residual blocks. - side_channels: Width of the precomputed conditioning information. - num_steps: Number of discrete diffusion steps. - num_layers: Number of residual blocks, at least one. - num_heads: Attention heads for each axis. - embedding_channels: Even diffusion and strategy embedding width, - at least four. - input_channels: Input width (two for conditional, one unconditional). - hidden_channels: Transformer feedforward width. - dropout: Transformer attention and feedforward dropout probability. - device: The device. - dtype: The parameter and output dtype. - """ - def __init__( self, channels: int, @@ -40,15 +22,14 @@ def __init__( num_steps: int, num_layers: int, num_heads: int, - embedding_channels: int = 128, + embedding_channels: int, input_channels: int = 2, - hidden_channels: int = 64, - dropout: float = 0.1, 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, @@ -59,77 +40,54 @@ def __init__( embedding_dim=embedding_channels, **factory_kwargs, ) - self.input_projection = Conv1d( - in_channels=input_channels, - out_channels=channels, - kernel_size=1, + self.input_projection = Linear( + input_channels, + channels, **factory_kwargs, ) - self.output_projection1 = Conv1d( - in_channels=channels, - out_channels=channels, - kernel_size=1, - **factory_kwargs, - ) - self.output_projection2 = Conv1d( - in_channels=channels, - out_channels=1, - kernel_size=1, - **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) + ] ) - for projection in (self.input_projection, self.output_projection1): + 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=projection.weight, mode="fan_out", nonlinearity="relu" + tensor=lin.weight, + mode="fan_out", + nonlinearity="relu", ) with torch.no_grad(): - projection.weight.mul_(0.5) + lin.weight.mul_(0.5) torch.nn.init.xavier_uniform_( - tensor=self.output_projection2.weight, gain=0.01 + tensor=self.output_projection2.weight, + gain=0.01, ) - self.residual_layers = ModuleList( - ResidualBlock( - channels=channels, - side_channels=side_channels, - embedding_channels=embedding_channels, - num_heads=num_heads, - hidden_channels=hidden_channels, - dropout=dropout, - **factory_kwargs, - ) - for _ in range(num_layers) - ) - self.norm_after_skip = LayerNorm(channels, **factory_kwargs) - self.norm_after_proj1 = LayerNorm(channels, **factory_kwargs) def forward( self, - x: Tensor, - side_info: Tensor, - step: Tensor, - strategy: Tensor, - ) -> Tensor: - """Predict noise for normalized time-series windows. - - Args: - x: Input with shape ``[B, I, K, L]``, where ``B`` is batch size, - ``I`` input channels, ``K`` features, and ``L`` time steps. - side_info: Conditioning with shape ``[B, S, K, L]``, where - ``S`` is ``side_channels``. - step: Integer diffusion indices in ``[0, num_steps)`` with - shape ``[B]`` or ``[1]``. - strategy: Integer masking-strategy indices in ``{0, 1}`` with - shape ``[B]`` or ``[1]``. Single-entry indices broadcast - across the batch independently for either embedding. + 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.movedim(1, -1)).relu() # [B, K, L, C] + x = x.movedim(-1, 1) # [B, C, K, L] - Returns: - Predicted noise with shape ``[B, K, L]`` in input feature/time - order. No denormalization or mask restoration is applied. - """ - batch, _, features, length = x.shape - x = self.input_projection(x.flatten(2)).relu() - x = x.reshape(batch, -1, features, length) diffusion_embedding = self.diffusion_embedding(step) strategy_embedding = self.strategy_embedding(strategy) + skips = [] for layer in self.residual_layers: x, skip = layer( @@ -139,10 +97,10 @@ def forward( strategy_embedding=strategy_embedding, ) skips.append(skip) - x = torch.stack(skips).sum(dim=0) / math.sqrt(len(skips)) - x = self.norm_after_skip(x.flatten(2).transpose(1, 2)) # [B, K * L, C] - x = self.output_projection1(x.transpose(1, 2)) - x = self.norm_after_proj1(x.transpose(1, 2)).transpose(1, 2) - return self.output_projection2(x.relu()).reshape( - batch, features, length - ) + + x = torch.stack(skips).sum(dim=0) # [B, C, K, L] + x = x / math.sqrt(len(skips)) + x = self.norm_after_skip(x.movedim(1, -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 index f18f310b6c..bed3cadf75 100644 --- a/test/models/kumo/timeseries/anomaly/test_denoiser.py +++ b/test/models/kumo/timeseries/anomaly/test_denoiser.py @@ -1,120 +1,41 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import math - import pytest import torch -import torch.nn.functional as F -from torch import Tensor from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser from sdm.testing import withCUDA -def reference_denoiser( - model: DiffusionDenoiser, - x: Tensor, - side: Tensor, - step: Tensor, - strategy: Tensor, -) -> Tensor: - """Compose the tested blocks with pointwise linear projections.""" - x = x.permute(0, 2, 3, 1) - projection = model.input_projection - x = F.linear(x, projection.weight.squeeze(-1), projection.bias).relu() - x = x.permute(0, 3, 1, 2) - diffusion = model.diffusion_embedding(step) - strategy_embedding = model.strategy_embedding(strategy) - skips = [] - for block in model.residual_layers: - x, skip = block(x, side, diffusion, strategy_embedding) - skips.append(skip.permute(0, 2, 3, 1)) - x = sum(skips) / math.sqrt(len(skips)) - x = model.norm_after_skip(x) - projection = model.output_projection1 - x = F.linear(x, projection.weight.squeeze(-1), projection.bias) - x = model.norm_after_proj1(x).relu() - projection = model.output_projection2 - x = F.linear(x, projection.weight.squeeze(-1), projection.bias) - return x.squeeze(-1) - - @withCUDA -@pytest.mark.parametrize( - ( - "input_channels", - "num_layers", - "features", - "length", - "step_batch", - "strategy_batch", - ), - [ - (2, 1, 3, 5, 1, 2), - (1, 3, 3, 5, 2, 1), - (2, 2, 1, 1, 2, 2), - ], -) -def test_denoiser( - device: torch.device, +@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, - num_layers: int, - step_batch: int, - strategy_batch: int, + device: torch.device, ) -> None: model = DiffusionDenoiser( channels=8, side_channels=6, num_steps=10, - num_layers=num_layers, + num_layers=2, num_heads=2, embedding_channels=12, input_channels=input_channels, device=device, - ).eval() + ) + x = torch.randn(2, input_channels, features, length, device=device) - side = torch.randn(2, 6, features, length, device=device) - step = torch.tensor([0, 9], device=device)[:step_batch] - strategy = torch.tensor([0, 1], device=device)[:strategy_batch] - with torch.inference_mode(): - actual = model(x, side, step, strategy) - expected = reference_denoiser(model, x, side, step, strategy) - assert actual.shape == (2, features, length) - assert actual.dtype == x.dtype - assert actual.device == device - assert actual.isfinite().all() - torch.testing.assert_close(actual, expected) + 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) -@withCUDA -def test_denoiser_grad(device: torch.device) -> None: - model = DiffusionDenoiser( - channels=8, - side_channels=6, - num_steps=10, - num_layers=2, - num_heads=2, - embedding_channels=12, - device=device, - dtype=torch.float64, - ).eval() - x = torch.randn( - 2, 2, 3, 5, device=device, dtype=torch.float64, requires_grad=True - ) - side = torch.randn( - 2, 6, 3, 5, device=device, dtype=torch.float64, requires_grad=True - ) - step = torch.tensor([0, 9], device=device) - strategy = torch.tensor([0, 1], device=device) - out = model(x, side, step, strategy) - assert out.dtype == torch.float64 - out.square().mean().backward() - for value in (x, side, *model.parameters()): - assert value.grad is not None - assert value.grad.isfinite().all() - for value in (x, side): - assert value.grad is not None - assert value.grad.abs().sum() > 0 + assert out.size() == (2, features, length) + assert out.dtype == x.dtype + assert out.device == device From 66d2fa0c5b1b03ea38b3f5c707841b4c33e50226 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:46:17 -0700 Subject: [PATCH 4/4] Use `permute` for channel moves in `DiffusionDenoiser` --- sdm/models/kumo/timeseries/anomaly/denoiser.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py index 371faa7776..f9e4931c6e 100644 --- a/sdm/models/kumo/timeseries/anomaly/denoiser.py +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -82,8 +82,8 @@ def forward( step: Tensor, # [B] or [1] strategy: Tensor, # [B] or [1] ) -> Tensor: # [B, K, L] - x = self.input_projection(x.movedim(1, -1)).relu() # [B, K, L, C] - x = x.movedim(-1, 1) # [B, C, 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) @@ -100,7 +100,7 @@ def forward( x = torch.stack(skips).sum(dim=0) # [B, C, K, L] x = x / math.sqrt(len(skips)) - x = self.norm_after_skip(x.movedim(1, -1)) # [B, K, L, C] + 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]