From 200f0d037b1db96bfee96c4ed269145e6bbb6cf4 Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Sat, 14 Mar 2026 17:29:20 +0800 Subject: [PATCH 1/5] standardize backbone API with forward_features and compute_features_output_shape --- test/test_models/test_backbone_api.py | 71 ++++++++++++++++ torch_ecg/models/cnn/darknet.py | 26 ++++-- torch_ecg/models/cnn/densenet.py | 38 +++++++++ torch_ecg/models/cnn/efficientnet.py | 40 ++++++--- torch_ecg/models/cnn/ho_resnet.py | 64 ++++++++++++++- torch_ecg/models/cnn/mobilenet.py | 114 ++++++++++++++++++++++++++ torch_ecg/models/cnn/multi_scopic.py | 38 +++++++++ torch_ecg/models/cnn/regnet.py | 40 ++++++++- torch_ecg/models/cnn/resnet.py | 38 +++++++++ torch_ecg/models/cnn/vgg.py | 40 ++++++++- torch_ecg/models/cnn/xception.py | 38 +++++++++ torch_ecg/utils/utils_nn.py | 24 ++++++ 12 files changed, 551 insertions(+), 20 deletions(-) create mode 100644 test/test_models/test_backbone_api.py diff --git a/test/test_models/test_backbone_api.py b/test/test_models/test_backbone_api.py new file mode 100644 index 00000000..78586505 --- /dev/null +++ b/test/test_models/test_backbone_api.py @@ -0,0 +1,71 @@ +""" +Unit tests for the standardized Backbone API. +""" + +from copy import deepcopy + +import pytest +import torch + +from torch_ecg.model_configs import ECG_CRNN_CONFIG +from torch_ecg.models.registry import BACKBONES + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +@pytest.mark.parametrize("backbone_name", BACKBONES.list_all()) +def test_backbone_api(backbone_name): + # Skip aliases to avoid redundant tests + if backbone_name != backbone_name.lower(): + return + + n_leads = 12 + batch_size = 2 + seq_len = 2000 + + # Get default config if available in ECG_CRNN_CONFIG + config = None + for k, v in ECG_CRNN_CONFIG.cnn.items(): + if k.lower() == backbone_name.lower(): + config = deepcopy(v) + break + + if config is None: + # Some backbones might not be in ECG_CRNN_CONFIG, skip for now + # or provide a minimal dummy config if known + pytest.skip(f"No default config found for backbone: {backbone_name}") + + try: + model = BACKBONES.build(backbone_name, in_channels=n_leads, **config).to(DEVICE) + except Exception as e: + pytest.fail(f"Failed to build backbone {backbone_name} with config {config}: {e}") + + model.eval() + inp = torch.randn(batch_size, n_leads, seq_len).to(DEVICE) + + # 1. Test forward_features existence + assert hasattr(model, "forward_features"), f"Backbone {backbone_name} missing forward_features" + + # 2. Test forward_features output shape + features = model.forward_features(inp) + assert features.ndim == 3, f"Backbone {backbone_name} forward_features should return 3D tensor, got {features.ndim}D" + assert features.shape[0] == batch_size + + # 3. Test compute_features_output_shape consistency + expected_shape = model.compute_features_output_shape(seq_len, batch_size) + assert ( + features.shape[1] == expected_shape[1] + ), f"Backbone {backbone_name} feature channels mismatch: {features.shape[1]} vs {expected_shape[1]}" + if expected_shape[2] is not None: + assert ( + features.shape[2] == expected_shape[2] + ), f"Backbone {backbone_name} feature seq_len mismatch: {features.shape[2]} vs {expected_shape[2]}" + + # 4. Test forward consistency (if model is pure feature extractor) + # For now, all current backbones in torch_ecg are pure feature extractors + out = model(inp) + assert torch.allclose(out, features), f"Backbone {backbone_name} forward and forward_features results differ" + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/torch_ecg/models/cnn/darknet.py b/torch_ecg/models/cnn/darknet.py index 3b970fbb..0d4a68d5 100644 --- a/torch_ecg/models/cnn/darknet.py +++ b/torch_ecg/models/cnn/darknet.py @@ -10,19 +10,15 @@ 5. Wang, C. Y., Bochkovskiy, A., & Liao, H. Y. M. (2020). Scaled-YOLOv4: Scaling Cross Stage Partial Network. arXiv preprint arXiv:2011.08036. """ -from typing import List +from typing import List, Optional, Sequence, Union -from torch import nn +from torch import Tensor, nn from ...models._nets import Conv_Bn_Activation, DownSample, GlobalContextBlock, NonLocalBlock, SEBlock # noqa: F401 from ...utils import CitationMixin, SizeMixin -__all__ = [ - "DarkNet", -] - -class DarkNet(nn.Sequential, SizeMixin, CitationMixin): +class DarkNet(SizeMixin, nn.Sequential, CitationMixin): """ """ __name__ = "DarkNet" @@ -32,6 +28,22 @@ def __init__(self, in_channels: int, **config) -> None: super().__init__() raise NotImplementedError + def compute_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the model.""" + raise NotImplementedError + + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features.""" + raise NotImplementedError + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features.""" + raise NotImplementedError + @property def doi(self) -> List[str]: return list(set(self.config.get("doi", []) + ["10.1109/CVPR.2016.91"])) diff --git a/torch_ecg/models/cnn/densenet.py b/torch_ecg/models/cnn/densenet.py index 5fa89995..d0329bee 100644 --- a/torch_ecg/models/cnn/densenet.py +++ b/torch_ecg/models/cnn/densenet.py @@ -760,6 +760,44 @@ def compute_output_shape( """Compute the output shape of the network.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels diff --git a/torch_ecg/models/cnn/efficientnet.py b/torch_ecg/models/cnn/efficientnet.py index a54b19b5..caf5bae4 100644 --- a/torch_ecg/models/cnn/efficientnet.py +++ b/torch_ecg/models/cnn/efficientnet.py @@ -9,17 +9,13 @@ """ -from typing import List +from typing import List, Optional, Sequence, Union -from torch import nn +from torch import Tensor, nn from ...models._nets import Conv_Bn_Activation, DownSample, GlobalContextBlock, NonLocalBlock, SEBlock # noqa: F401 from ...utils import CitationMixin, SizeMixin -__all__ = [ - "EfficientNet", -] - class EfficientNet(nn.Module, SizeMixin, CitationMixin): """ @@ -38,10 +34,22 @@ def __init__(self, in_channels: int, **config) -> None: super().__init__() raise NotImplementedError - def forward(self): + def forward(self, input: Tensor) -> Tensor: + raise NotImplementedError + + def compute_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + raise NotImplementedError + + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features.""" raise NotImplementedError - def compute_output_shape(self): + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features.""" raise NotImplementedError @property @@ -67,8 +75,20 @@ def __init__(self, in_channels: int, **config) -> None: super().__init__() raise NotImplementedError - def forward(self): + def forward(self, input: Tensor) -> Tensor: + raise NotImplementedError + + def compute_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + raise NotImplementedError + + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features.""" raise NotImplementedError - def compute_output_shape(self): + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features.""" raise NotImplementedError diff --git a/torch_ecg/models/cnn/ho_resnet.py b/torch_ecg/models/cnn/ho_resnet.py index 8f0f485c..deb5ab95 100644 --- a/torch_ecg/models/cnn/ho_resnet.py +++ b/torch_ecg/models/cnn/ho_resnet.py @@ -7,7 +7,9 @@ """ -from torch import nn +from typing import Optional, Sequence, Union + +from torch import Tensor, nn from ...cfg import CFG # noqa: F401 from ...models._nets import ( # noqa: F401 @@ -33,6 +35,26 @@ class MidPointResNet(nn.Module, SizeMixin, CitationMixin): def __init__(self, in_channels: int, **config) -> None: """ """ + super().__init__() + raise NotImplementedError + + def forward(self, input: Tensor) -> Tensor: + raise NotImplementedError + + def compute_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the model.""" + raise NotImplementedError + + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features.""" + raise NotImplementedError + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features.""" raise NotImplementedError @@ -41,6 +63,26 @@ class RK4ResNet(nn.Module, SizeMixin, CitationMixin): def __init__(self, in_channels: int, **config) -> None: """ """ + super().__init__() + raise NotImplementedError + + def forward(self, input: Tensor) -> Tensor: + raise NotImplementedError + + def compute_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the model.""" + raise NotImplementedError + + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features.""" + raise NotImplementedError + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features.""" raise NotImplementedError @@ -49,4 +91,24 @@ class RK8ResNet(nn.Module, SizeMixin, CitationMixin): def __init__(self, in_channels: int, **config) -> None: """ """ + super().__init__() + raise NotImplementedError + + def forward(self, input: Tensor) -> Tensor: + raise NotImplementedError + + def compute_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the model.""" + raise NotImplementedError + + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features.""" + raise NotImplementedError + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features.""" raise NotImplementedError diff --git a/torch_ecg/models/cnn/mobilenet.py b/torch_ecg/models/cnn/mobilenet.py index cd6e476b..f50a6ba8 100644 --- a/torch_ecg/models/cnn/mobilenet.py +++ b/torch_ecg/models/cnn/mobilenet.py @@ -414,6 +414,44 @@ def compute_output_shape( _, _, _seq_len = output_shape return output_shape + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels @@ -829,6 +867,44 @@ def compute_output_shape( """Compute the output shape of the model.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels @@ -1253,6 +1329,44 @@ def compute_output_shape( """Compute the output shape of the model.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels diff --git a/torch_ecg/models/cnn/multi_scopic.py b/torch_ecg/models/cnn/multi_scopic.py index e999a937..4be1144e 100644 --- a/torch_ecg/models/cnn/multi_scopic.py +++ b/torch_ecg/models/cnn/multi_scopic.py @@ -480,6 +480,44 @@ def compute_output_shape( output_shape = (batch_size, out_channels, _seq_len) return output_shape + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + def assign_weights_lead_wise(self, other: "MultiScopicCNN", indices: Sequence[int]) -> None: """Assign weights to the `other` :class:`MultiScopicCNN` module in the lead-wise manner diff --git a/torch_ecg/models/cnn/regnet.py b/torch_ecg/models/cnn/regnet.py index 0f3cab31..650989f9 100644 --- a/torch_ecg/models/cnn/regnet.py +++ b/torch_ecg/models/cnn/regnet.py @@ -11,7 +11,7 @@ from typing import List, Optional, Sequence, Union import torch -from torch import nn +from torch import Tensor, nn from ...cfg import CFG from ...models._nets import Conv_Bn_Activation, DownSample, SpaceToDepth @@ -512,6 +512,44 @@ def compute_output_shape( """Compute the output shape of the network.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels diff --git a/torch_ecg/models/cnn/resnet.py b/torch_ecg/models/cnn/resnet.py index 314ed107..b7b72dd7 100644 --- a/torch_ecg/models/cnn/resnet.py +++ b/torch_ecg/models/cnn/resnet.py @@ -858,6 +858,44 @@ def compute_output_shape( """Compute the output shape of the model.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels diff --git a/torch_ecg/models/cnn/vgg.py b/torch_ecg/models/cnn/vgg.py index 9157ef15..00f2069b 100644 --- a/torch_ecg/models/cnn/vgg.py +++ b/torch_ecg/models/cnn/vgg.py @@ -4,7 +4,7 @@ from copy import deepcopy from typing import List, Optional, Sequence, Union -from torch import nn +from torch import Tensor, nn from ...cfg import CFG from ...models._nets import Conv_Bn_Activation @@ -189,6 +189,44 @@ def compute_output_shape( """Compute the output shape of the module.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels diff --git a/torch_ecg/models/cnn/xception.py b/torch_ecg/models/cnn/xception.py index e49e1dab..85308e6b 100644 --- a/torch_ecg/models/cnn/xception.py +++ b/torch_ecg/models/cnn/xception.py @@ -888,6 +888,44 @@ def compute_output_shape( """Compute the output shape the model.""" return compute_sequential_output_shape(self, seq_len, batch_size) + def forward_features(self, input: Tensor) -> Tensor: + """Forward pass of the model to extract features. + + Parameters + ---------- + input : torch.Tensor + Input signal tensor, + of shape ``(batch_size, channels, seq_len)``. + + Returns + ------- + features : torch.Tensor + Feature map tensor, + of shape ``(batch_size, channels, seq_len)``. + + """ + return self.forward(input) + + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + @property def in_channels(self) -> int: return self.__in_channels diff --git a/torch_ecg/utils/utils_nn.py b/torch_ecg/utils/utils_nn.py index 63f84e53..8847d169 100644 --- a/torch_ecg/utils/utils_nn.py +++ b/torch_ecg/utils/utils_nn.py @@ -1073,6 +1073,30 @@ def dtype_(self) -> str: def device_(self) -> str: return str(self.device) + def compute_features_output_shape( + self, seq_len: Optional[int] = None, batch_size: Optional[int] = None + ) -> Sequence[Union[int, None]]: + """Compute the output shape of the features. + + By default, this is the same as the output shape of the model. + For backbones with pooling and classification heads, this should be + overridden to return the shape of the features before global pooling. + + Parameters + ---------- + seq_len : int, optional + Length of the input signal tensor. + batch_size : int, optional + Batch size of the input signal tensor. + + Returns + ------- + output_shape : sequence + Output shape of the features. + + """ + return self.compute_output_shape(seq_len, batch_size) + def make_safe_globals(obj: Any, remove_paths: bool = True) -> Any: """Make a dictionary or a dictionary-like object safe for serialization. From cbc7a352190560c7fa0a6eb2383548a70ad5e1c5 Mon Sep 17 00:00:00 2001 From: WEN Hao <8778305+wenh06@users.noreply.github.com> Date: Sat, 14 Mar 2026 21:53:15 +0800 Subject: [PATCH 2/5] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- test/test_models/test_backbone_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_models/test_backbone_api.py b/test/test_models/test_backbone_api.py index 78586505..32c98e4c 100644 --- a/test/test_models/test_backbone_api.py +++ b/test/test_models/test_backbone_api.py @@ -17,7 +17,7 @@ def test_backbone_api(backbone_name): # Skip aliases to avoid redundant tests if backbone_name != backbone_name.lower(): - return + pytest.skip(f"Skipping alias backbone name: {backbone_name}") n_leads = 12 batch_size = 2 From 510614043ed95e31b21fdba3808b9ea16eba593d Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Sun, 15 Mar 2026 11:59:28 +0800 Subject: [PATCH 3/5] standardize backbone API with forward_features and compute_features_output_shape --- CHANGELOG.rst | 15 ++++++++ test/test_models/test_backbone_api.py | 55 ++++++++++++++++----------- torch_ecg/utils/utils_nn.py | 14 +++---- 3 files changed, 54 insertions(+), 30 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a4d9a63a..d0fbf3b6 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,9 +14,17 @@ Versioning `__. Added ~~~~~ +- Implement a generic `Registry` class and establish registries for backbones, models, + attention layers, preprocessors, augmenters, optimizers, schedulers, and losses. +- Add `forward_features` and `compute_features_output_shape` methods to all CNN + backbones to provide a standardized API for SSL and feature extraction. + Changed ~~~~~~~ +- Refactor `ECG_CRNN`, `PreprocManager`, `AugmenterManager`, and `BaseTrainer` to + utilize the new registry system for dynamic component construction and decoupling. +- Enhance `SizeMixin` to support static shape inference for feature maps. - Make the function `remove_spikes_naive` in `torch_ecg.utils.utils_signal` support 2D and 3D input signals. - Use `save_file` and `load_file` from the `safetensors` package for saving @@ -36,6 +44,13 @@ Removed Fixed ~~~~~ +- Robustly handle dimension inference and initialization in `ECG_CRNN` models, + especially for cases with `None` or `Identity` modules. +- Address several compatibility issues for Python 3.13, including docstring + indentation and `NaN` comparisons in dataclasses. +- Improve error handling and encoding robustness in `CitationMixin` when + reading cache files. +- Resolve CodeQL warnings regarding incomplete URL substring sanitization in tests. - Correctly update the `_df_metadata` attribute of the `PTBXL` database reader classes after filtering records. - Enhance the `save` method of the `torch_ecg.utils.utils_nn.CkptMixin` class: diff --git a/test/test_models/test_backbone_api.py b/test/test_models/test_backbone_api.py index 78586505..7ab4fcba 100644 --- a/test/test_models/test_backbone_api.py +++ b/test/test_models/test_backbone_api.py @@ -12,33 +12,44 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") - -@pytest.mark.parametrize("backbone_name", BACKBONES.list_all()) -def test_backbone_api(backbone_name): - # Skip aliases to avoid redundant tests - if backbone_name != backbone_name.lower(): - return - +# Extract all valid backbone configurations from the central config +BACKBONE_CONFIGS = [] +for config_key, config_val in ECG_CRNN_CONFIG.cnn.items(): + if not isinstance(config_val, dict): + continue + + # 1. Try to get name from the config dict + # 2. If not in dict, use the key if it's in the registry + # 3. If key contains a known registry name (e.g. resnet_nature_comm), extract the base name + backbone_name = config_val.get("name") + if backbone_name is None: + if config_key in BACKBONES: + backbone_name = config_key + else: + # Check if config_key contains any registered name as a prefix + for registered_name in BACKBONES.list_all(): + if config_key.startswith(registered_name): + backbone_name = registered_name + break + + if backbone_name: + BACKBONE_CONFIGS.append((backbone_name, config_key, config_val)) + + +@pytest.mark.parametrize("backbone_name, config_key, config", BACKBONE_CONFIGS) +def test_backbone_api(backbone_name, config_key, config): n_leads = 12 batch_size = 2 seq_len = 2000 - # Get default config if available in ECG_CRNN_CONFIG - config = None - for k, v in ECG_CRNN_CONFIG.cnn.items(): - if k.lower() == backbone_name.lower(): - config = deepcopy(v) - break - - if config is None: - # Some backbones might not be in ECG_CRNN_CONFIG, skip for now - # or provide a minimal dummy config if known - pytest.skip(f"No default config found for backbone: {backbone_name}") - + # Skip models that are not implemented yet to avoid noisy test failures + # These will be implemented in Phase 1.5 try: - model = BACKBONES.build(backbone_name, in_channels=n_leads, **config).to(DEVICE) + model = BACKBONES.build(backbone_name, in_channels=n_leads, **deepcopy(config)).to(DEVICE) + except NotImplementedError: + pytest.skip(f"Backbone {backbone_name} (config: {config_key}) is not implemented yet.") except Exception as e: - pytest.fail(f"Failed to build backbone {backbone_name} with config {config}: {e}") + pytest.fail(f"Failed to build backbone {backbone_name} with config {config_key}: {e}") model.eval() inp = torch.randn(batch_size, n_leads, seq_len).to(DEVICE) @@ -52,6 +63,7 @@ def test_backbone_api(backbone_name): assert features.shape[0] == batch_size # 3. Test compute_features_output_shape consistency + # All Backbones in torch_ecg follow (seq_len, batch_size) signature expected_shape = model.compute_features_output_shape(seq_len, batch_size) assert ( features.shape[1] == expected_shape[1] @@ -62,7 +74,6 @@ def test_backbone_api(backbone_name): ), f"Backbone {backbone_name} feature seq_len mismatch: {features.shape[2]} vs {expected_shape[2]}" # 4. Test forward consistency (if model is pure feature extractor) - # For now, all current backbones in torch_ecg are pure feature extractors out = model(inp) assert torch.allclose(out, features), f"Backbone {backbone_name} forward and forward_features results differ" diff --git a/torch_ecg/utils/utils_nn.py b/torch_ecg/utils/utils_nn.py index 8847d169..d01ec960 100644 --- a/torch_ecg/utils/utils_nn.py +++ b/torch_ecg/utils/utils_nn.py @@ -1073,9 +1073,7 @@ def dtype_(self) -> str: def device_(self) -> str: return str(self.device) - def compute_features_output_shape( - self, seq_len: Optional[int] = None, batch_size: Optional[int] = None - ) -> Sequence[Union[int, None]]: + def compute_features_output_shape(self, *args: Any, **kwargs: Any) -> Sequence[Union[int, None]]: """Compute the output shape of the features. By default, this is the same as the output shape of the model. @@ -1084,10 +1082,10 @@ def compute_features_output_shape( Parameters ---------- - seq_len : int, optional - Length of the input signal tensor. - batch_size : int, optional - Batch size of the input signal tensor. + *args : Any + Positional arguments passed to `compute_output_shape`. + **kwargs : Any + Keyword arguments passed to `compute_output_shape`. Returns ------- @@ -1095,7 +1093,7 @@ def compute_features_output_shape( Output shape of the features. """ - return self.compute_output_shape(seq_len, batch_size) + return self.compute_output_shape(*args, **kwargs) def make_safe_globals(obj: Any, remove_paths: bool = True) -> Any: From 9c2201f45bb94a9b4d735bf1c26185e9b25f0357 Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Tue, 17 Mar 2026 21:18:19 +0800 Subject: [PATCH 4/5] update test coverage settings --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1be2b9df..23b4fd33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -158,6 +158,10 @@ omit = [ "torch_ecg/models/grad_cam.py", # temporarily ignore torch_ecg/ssl since it's not implemented "torch_ecg/ssl/*", + # temporarily ignore models that are not implemented completely + "torch_ecg/models/cnn/darknet.py", + "torch_ecg/models/cnn/efficientnet.py", + "torch_ecg/models/cnn/ho_resnet.py", ] exclude_also = [ "raise NotImplementedError", From b177990da16f17afd9ee608a38aad6e0bbfde573 Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Tue, 17 Mar 2026 21:18:58 +0800 Subject: [PATCH 5/5] update github action configs --- .github/workflows/auto-label.yml | 15 +++++++++++++++ .github/workflows/publish.yml | 32 +++++++++++++++++++++++++++++++- 2 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/auto-label.yml diff --git a/.github/workflows/auto-label.yml b/.github/workflows/auto-label.yml new file mode 100644 index 00000000..639f3264 --- /dev/null +++ b/.github/workflows/auto-label.yml @@ -0,0 +1,15 @@ +name: Auto Label + +on: + pull_request: + +jobs: + label: + runs-on: ubuntu-latest + steps: + - name: Label pre-commit PR + if: github.actor == 'pre-commit-ci[bot]' + uses: actions-ecosystem/action-add-labels@v1 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + labels: pre-commit diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 110f6bb8..ea0739d3 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -87,9 +87,39 @@ jobs: uses: actions/checkout@v4 - name: Build Changelog id: changelog - uses: mikepenz/release-changelog-builder-action@v4 + uses: mikepenz/release-changelog-builder-action@v6 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + configurationJson: | + { + "template": "# Changelog\n\n{{CHANGELOG}}\n\n
\n📦 Other changes\n\n{{UNCATEGORIZED}}\n\n
", + "categories": [ + { + "title": "🚀 Features", + "labels": ["feature","enhancement"] + }, + { + "title": "🐛 Bug Fixes", + "labels": ["bug","fix"] + }, + { + "title": "📚 Documentation", + "labels": ["documentation","docs"] + }, + { + "title": "⬆️ Dependencies", + "labels": ["dependencies"] + }, + { + "title": "🧰 Maintenance", + "labels": ["chore","ci"] + } + ], + "exclude_labels": [ + "pre-commit" + ] + } - name: Create Release id: create_release uses: softprops/action-gh-release@v2