From e82768274a6945de1b4369eea7245785d0e9dcb7 Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Thu, 12 Mar 2026 18:15:43 +0800 Subject: [PATCH 1/4] implement registries: the initial draft --- .github/workflows/test-databases.yml | 2 +- .github/workflows/test-models.yml | 2 +- .pre-commit-config.yaml | 2 +- torch_ecg/augmenters/__init__.py | 3 + torch_ecg/augmenters/augmenter_manager.py | 132 ++++++------------- torch_ecg/augmenters/baseline_wander.py | 3 + torch_ecg/augmenters/cutmix.py | 3 + torch_ecg/augmenters/label_smooth.py | 3 + torch_ecg/augmenters/mixup.py | 3 + torch_ecg/augmenters/random_flip.py | 3 + torch_ecg/augmenters/random_masking.py | 3 + torch_ecg/augmenters/random_renormalize.py | 3 + torch_ecg/augmenters/registry.py | 11 ++ torch_ecg/augmenters/stretch_compress.py | 3 + torch_ecg/components/registry.py | 15 +++ torch_ecg/components/trainer.py | 95 ++++++++------ torch_ecg/models/__init__.py | 4 + torch_ecg/models/cnn/densenet.py | 4 + torch_ecg/models/cnn/mobilenet.py | 10 ++ torch_ecg/models/cnn/multi_scopic.py | 4 + torch_ecg/models/cnn/regnet.py | 3 + torch_ecg/models/cnn/resnet.py | 4 + torch_ecg/models/cnn/vgg.py | 3 + torch_ecg/models/cnn/xception.py | 3 + torch_ecg/models/ecg_crnn.py | 80 +++++------- torch_ecg/models/ecg_seq_lab_net.py | 2 + torch_ecg/models/loss.py | 21 ++- torch_ecg/models/registry.py | 13 ++ torch_ecg/models/rr_lstm.py | 2 + torch_ecg/models/unets/ecg_subtract_unet.py | 2 + torch_ecg/models/unets/ecg_unet.py | 2 + torch_ecg/preprocessors/__init__.py | 3 + torch_ecg/preprocessors/bandpass.py | 3 + torch_ecg/preprocessors/baseline_remove.py | 3 + torch_ecg/preprocessors/normalize.py | 9 ++ torch_ecg/preprocessors/preproc_manager.py | 136 ++++++------------- torch_ecg/preprocessors/registry.py | 11 ++ torch_ecg/preprocessors/resample.py | 3 + torch_ecg/utils/__init__.py | 3 + torch_ecg/utils/registry.py | 137 ++++++++++++++++++++ 40 files changed, 456 insertions(+), 295 deletions(-) create mode 100644 torch_ecg/augmenters/registry.py create mode 100644 torch_ecg/components/registry.py create mode 100644 torch_ecg/models/registry.py create mode 100644 torch_ecg/preprocessors/registry.py create mode 100644 torch_ecg/utils/registry.py diff --git a/.github/workflows/test-databases.yml b/.github/workflows/test-databases.yml index c99c4edc..b78076e0 100644 --- a/.github/workflows/test-databases.yml +++ b/.github/workflows/test-databases.yml @@ -18,7 +18,7 @@ jobs: runs-on: ubuntu-22.04 strategy: matrix: - python-version: ["3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.12", "3.13"] fail-fast: false steps: diff --git a/.github/workflows/test-models.yml b/.github/workflows/test-models.yml index 9e885673..d804832e 100644 --- a/.github/workflows/test-models.yml +++ b/.github/workflows/test-models.yml @@ -18,7 +18,7 @@ jobs: runs-on: ubuntu-22.04 strategy: matrix: - python-version: ["3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.12", "3.13"] fail-fast: false steps: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bc43ceeb..7b69b57a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,7 +18,7 @@ repos: hooks: - id: flake8 additional_dependencies: [pycodestyle>=2.11.0] - args: [--max-line-length=128, '--exclude=./.*,build,dist', '--ignore=E501,W503,E203,F841,E231,W604,E226', --count, --statistics, --show-source] + args: [--max-line-length=128, '--exclude=./.*,build,dist', '--ignore=E501,W503,E203,F841,E231,W604,E226,E225', --count, --statistics, --show-source] - repo: https://github.com/pycqa/isort rev: 8.0.1 hooks: diff --git a/torch_ecg/augmenters/__init__.py b/torch_ecg/augmenters/__init__.py index 2d307ae7..f0610149 100644 --- a/torch_ecg/augmenters/__init__.py +++ b/torch_ecg/augmenters/__init__.py @@ -15,6 +15,7 @@ :toctree: generated/ :recursive: + AUGMENTERS AugmenterManager Augmenter BaselineWanderAugmenter @@ -38,9 +39,11 @@ from .random_flip import RandomFlip from .random_masking import RandomMasking from .random_renormalize import RandomRenormalize +from .registry import AUGMENTERS from .stretch_compress import StretchCompress, StretchCompressOffline __all__ = [ + "AUGMENTERS", "Augmenter", "BaselineWanderAugmenter", "CutMix", diff --git a/torch_ecg/augmenters/augmenter_manager.py b/torch_ecg/augmenters/augmenter_manager.py index 3c5e7e4a..197522d2 100644 --- a/torch_ecg/augmenters/augmenter_manager.py +++ b/torch_ecg/augmenters/augmenter_manager.py @@ -9,13 +9,7 @@ from ..utils.misc import add_docstring, default_class_repr from .base import Augmenter, _augmenter_forward_doc -from .baseline_wander import BaselineWanderAugmenter -from .label_smooth import LabelSmooth -from .mixup import Mixup -from .random_flip import RandomFlip -from .random_masking import RandomMasking -from .random_renormalize import RandomRenormalize -from .stretch_compress import StretchCompress +from .registry import AUGMENTERS __all__ = [ "AugmenterManager", @@ -65,84 +59,34 @@ class AugmenterManager(torch.nn.Module): def __init__(self, *augs: Optional[Tuple[Augmenter, ...]], random: bool = False) -> None: super().__init__() self.random = random - self._augmenters = list(augs) + self._augmenters = torch.nn.ModuleList(list(augs)) - def _add_baseline_wander(self, **config: dict) -> None: - """Add the baseline wander augmenter to the manager. + def add_(self, aug: Union[Augmenter, str, dict], pos: int = -1, **kwargs: Any) -> None: + """Add an augmenter to the manager. Parameters ---------- - **config : dict - The configuration for the baseline wander augmenter. + aug : Augmenter or str or dict + The augmenter to be added. + If it's a string or a dict, it will be built using the registry. + pos : int, default -1 + The position to insert the augmenter. + Should be >= -1, with -1 being the indicator of the end. + **kwargs : Any + Additional keyword arguments for building the augmenter. """ - self._augmenters.append(BaselineWanderAugmenter(**config)) - - def _add_label_smooth(self, **config: dict) -> None: - """Add the label smooth augmenter to the manager. - - Parameters - ---------- - **config : dict - The configuration for the label smooth augmenter. - - """ - self._augmenters.append(LabelSmooth(**config)) - - def _add_mixup(self, **config: dict) -> None: - """Add the mixup augmenter to the manager. - - Parameters - ---------- - **config : dict - The configuration for the mixup augmenter. - - """ - self._augmenters.append(Mixup(**config)) - - def _add_random_flip(self, **config: dict) -> None: - """Add the random flip augmenter to the manager. - - Parameters - ---------- - **config : dict - The configuration for the random flip augmenter. - - """ - self._augmenters.append(RandomFlip(**config)) - - def _add_random_masking(self, **config: dict) -> None: - """Add the random masking augmenter to the manager. - - Parameters - ---------- - **config : dict - The configuration for the random masking augmenter. - - """ - self._augmenters.append(RandomMasking(**config)) - - def _add_random_renormalize(self, **config: dict) -> None: - """Add the random renormalize augmenter to the manager. - - Parameters - ---------- - **config : dict - The configuration for the random renormalize augmenter. - - """ - self._augmenters.append(RandomRenormalize(**config)) - - def _add_stretch_compress(self, **config: dict) -> None: - """Add the stretch compress augmenter to the manager. - - Parameters - ---------- - **config : dict - The configuration for the stretch compress augmenter. - - """ - self._augmenters.append(StretchCompress(**config)) + if isinstance(aug, (str, dict)): + aug = AUGMENTERS.build(aug, **kwargs) + assert isinstance(aug, Augmenter) + assert aug.__class__.__name__ not in [ + a.__class__.__name__ for a in self.augmenters + ], f"Augmenter {aug.__class__.__name__} already exists." + assert isinstance(pos, int) and pos >= -1, f"pos must be an integer >= -1, but got {pos}." + if pos == -1: + self._augmenters.append(aug) + else: + self._augmenters.insert(pos, aug) @add_docstring( _augmenter_forward_doc.replace( @@ -168,7 +112,7 @@ def forward( return (sig, label, *extra_tensors) @property - def augmenters(self) -> List[Augmenter]: + def augmenters(self) -> torch.nn.ModuleList: """The list of augmenters in the manager.""" return self._augmenters @@ -199,23 +143,22 @@ def from_config(cls, config: dict) -> "AugmenterManager": """ am = cls(random=config.get("random", False)) - _mapping = { - "baseline_wander": am._add_baseline_wander, - "label_smooth": am._add_label_smooth, - "mixup": am._add_mixup, - "random_flip": am._add_random_flip, - "random_masking": am._add_random_masking, - "random_renormalize": am._add_random_renormalize, - "stretch_compress": am._add_stretch_compress, - } for aug_name, aug_config in config.items(): if aug_name in [ "fs", "random", ]: continue - elif aug_name in _mapping and isinstance(aug_config, dict): - _mapping[aug_name](fs=config["fs"], **aug_config) + if aug_name in AUGMENTERS or aug_name in [ + "".join([w.capitalize() for w in k.split("_")]) for k in AUGMENTERS.list_all() + ]: + if isinstance(aug_config, dict): + # add default fs from config if not specified in aug_config + if "fs" not in aug_config and "fs" in config: + aug_config["fs"] = config["fs"] + am.add_(aug_name, **aug_config) + else: + am.add_(aug_name) else: # just ignore the other items pass @@ -236,9 +179,10 @@ def rearrange(self, new_ordering: List[str]) -> None: "The augmenters are applied in random order, " "rearranging the augmenters will not take effect.", RuntimeWarning, ) + # TODO: use a more robust way to map names to classes _mapping = { # built-in augmenters "".join([w.capitalize() for w in k.split("_")]): k - for k in "label_smooth,mixup,random_flip,random_masking,random_renormalize,stretch_compress".split(",") + for k in "label_smooth,mixup,random_flip,random_masking,random_renormalize,stretch_compress,cutmix".split(",") } _mapping.update({"BaselineWanderAugmenter": "baseline_wander"}) _mapping.update({v: k for k, v in _mapping.items()}) @@ -250,4 +194,6 @@ def rearrange(self, new_ordering: List[str]) -> None: assert len(new_ordering) == len(set(new_ordering)), "Duplicate augmenter names." assert len(new_ordering) == len(self._augmenters), "Number of augmenters mismatch." - self._augmenters.sort(key=lambda aug: new_ordering.index(_mapping[aug.__class__.__name__])) + augs = list(self._augmenters) + augs.sort(key=lambda aug: new_ordering.index(_mapping[aug.__class__.__name__])) + self._augmenters = torch.nn.ModuleList(augs) diff --git a/torch_ecg/augmenters/baseline_wander.py b/torch_ecg/augmenters/baseline_wander.py index 57e48dcb..292e7b14 100644 --- a/torch_ecg/augmenters/baseline_wander.py +++ b/torch_ecg/augmenters/baseline_wander.py @@ -13,12 +13,15 @@ from ..cfg import DEFAULTS from ..utils.utils_signal import get_ampl from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "BaselineWanderAugmenter", ] +@AUGMENTERS.register(name="baseline_wander") +@AUGMENTERS.register() class BaselineWanderAugmenter(Augmenter): """Generate baseline wander composed of sinusoidal and Gaussian noise. diff --git a/torch_ecg/augmenters/cutmix.py b/torch_ecg/augmenters/cutmix.py index 60bd98f2..25a744e3 100644 --- a/torch_ecg/augmenters/cutmix.py +++ b/torch_ecg/augmenters/cutmix.py @@ -11,12 +11,15 @@ from ..cfg import DEFAULTS from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "CutMix", ] +@AUGMENTERS.register(name="cutmix") +@AUGMENTERS.register() class CutMix(Augmenter): """CutMix augmentation. diff --git a/torch_ecg/augmenters/label_smooth.py b/torch_ecg/augmenters/label_smooth.py index 5c6ae053..29ab0dea 100644 --- a/torch_ecg/augmenters/label_smooth.py +++ b/torch_ecg/augmenters/label_smooth.py @@ -6,12 +6,15 @@ from torch import Tensor from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "LabelSmooth", ] +@AUGMENTERS.register(name="label_smooth") +@AUGMENTERS.register() class LabelSmooth(Augmenter): """Label smoothing augmentation. diff --git a/torch_ecg/augmenters/mixup.py b/torch_ecg/augmenters/mixup.py index 4b60ee2c..f73083f4 100644 --- a/torch_ecg/augmenters/mixup.py +++ b/torch_ecg/augmenters/mixup.py @@ -11,12 +11,15 @@ from ..cfg import DEFAULTS from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "Mixup", ] +@AUGMENTERS.register(name="mixup") +@AUGMENTERS.register() class Mixup(Augmenter): """Mixup augmentor. diff --git a/torch_ecg/augmenters/random_flip.py b/torch_ecg/augmenters/random_flip.py index 7ceee5c4..0401ee47 100644 --- a/torch_ecg/augmenters/random_flip.py +++ b/torch_ecg/augmenters/random_flip.py @@ -8,12 +8,15 @@ from torch import Tensor from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "RandomFlip", ] +@AUGMENTERS.register(name="random_flip") +@AUGMENTERS.register() class RandomFlip(Augmenter): """Randomly flip the ECGs along the voltage axis. diff --git a/torch_ecg/augmenters/random_masking.py b/torch_ecg/augmenters/random_masking.py index 7ecb1582..768553e9 100644 --- a/torch_ecg/augmenters/random_masking.py +++ b/torch_ecg/augmenters/random_masking.py @@ -9,12 +9,15 @@ from torch import Tensor from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "RandomMasking", ] +@AUGMENTERS.register(name="random_masking") +@AUGMENTERS.register() class RandomMasking(Augmenter): """Randomly mask ECGs with a probability. diff --git a/torch_ecg/augmenters/random_renormalize.py b/torch_ecg/augmenters/random_renormalize.py index 5715a66c..69fad112 100644 --- a/torch_ecg/augmenters/random_renormalize.py +++ b/torch_ecg/augmenters/random_renormalize.py @@ -9,12 +9,15 @@ from ..cfg import DEFAULTS from ..utils.utils_signal_t import normalize as normalize_t from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "RandomRenormalize", ] +@AUGMENTERS.register(name="random_renormalize") +@AUGMENTERS.register() class RandomRenormalize(Augmenter): """Randomly re-normalize the ECG tensor, using the Z-score normalization method. diff --git a/torch_ecg/augmenters/registry.py b/torch_ecg/augmenters/registry.py new file mode 100644 index 00000000..d3e68c50 --- /dev/null +++ b/torch_ecg/augmenters/registry.py @@ -0,0 +1,11 @@ +""" +Registry for augmenters. +""" + +from ..utils.registry import Registry + +__all__ = [ + "AUGMENTERS", +] + +AUGMENTERS = Registry("augmenters") diff --git a/torch_ecg/augmenters/stretch_compress.py b/torch_ecg/augmenters/stretch_compress.py index 1337dc0e..6024b655 100644 --- a/torch_ecg/augmenters/stretch_compress.py +++ b/torch_ecg/augmenters/stretch_compress.py @@ -14,6 +14,7 @@ from ..cfg import DEFAULTS from ..utils.misc import ReprMixin, add_docstring from .base import Augmenter +from .registry import AUGMENTERS __all__ = [ "StretchCompress", @@ -21,6 +22,8 @@ ] +@AUGMENTERS.register(name="stretch_compress") +@AUGMENTERS.register() class StretchCompress(Augmenter): """Stretch-or-compress augmenter on ECG tensors. diff --git a/torch_ecg/components/registry.py b/torch_ecg/components/registry.py new file mode 100644 index 00000000..b4ae492d --- /dev/null +++ b/torch_ecg/components/registry.py @@ -0,0 +1,15 @@ +""" +Registries for components (optimizers, schedulers, losses, etc.). +""" + +from ..utils.registry import Registry + +__all__ = [ + "OPTIMIZERS", + "SCHEDULERS", + "LOSSES", +] + +OPTIMIZERS = Registry("optimizers") +SCHEDULERS = Registry("schedulers") +LOSSES = Registry("losses") diff --git a/torch_ecg/components/trainer.py b/torch_ecg/components/trainer.py index 249b4f42..f98d8ae0 100644 --- a/torch_ecg/components/trainer.py +++ b/torch_ecg/components/trainer.py @@ -24,6 +24,7 @@ from ..utils.misc import ReprMixin, dict_to_str, dicts_equal, get_date_str, get_kwargs from ..utils.utils_nn import default_collate_fn, make_safe_globals from .loggers import LoggerManager +from .registry import OPTIMIZERS, SCHEDULERS __all__ = [ "BaseTrainer", @@ -614,74 +615,86 @@ def n_val(self) -> int: def _setup_optimizer(self) -> None: """Setup the optimizer.""" - if self.train_config.optimizer.lower() == "adam": # type: ignore - optimizer_kwargs = get_kwargs(optim.Adam) + opt_name = self.train_config.optimizer.lower() + if opt_name == "adamw_amsgrad": + opt_name = "adamw" + self.train_config.amsgrad = True + elif opt_name == "adam_amsgrad": + opt_name = "adam" + self.train_config.amsgrad = True + + if opt_name not in OPTIMIZERS: + # try to find in torch.optim + for name in dir(optim): + if name.lower() == opt_name: + OPTIMIZERS.register(name=opt_name)(getattr(optim, name)) + break + + if opt_name in OPTIMIZERS: + opt_cls = OPTIMIZERS.get(opt_name) + optimizer_kwargs = get_kwargs(opt_cls) optimizer_kwargs.update({k: self.train_config.get(k, v) for k, v in optimizer_kwargs.items()}) optimizer_kwargs.update(dict(lr=self.lr)) - self.optimizer = optim.Adam( - params=self.model.parameters(), - **optimizer_kwargs, - ) - elif self.train_config.optimizer.lower() in ["adamw", "adamw_amsgrad"]: # type: ignore - optimizer_kwargs = get_kwargs(optim.AdamW) - optimizer_kwargs.update({k: self.train_config.get(k, v) for k, v in optimizer_kwargs.items()}) - optimizer_kwargs.update( - dict( - lr=self.lr, - amsgrad=self.train_config.optimizer.lower().endswith("amsgrad"), # type: ignore - ) - ) - self.optimizer = optim.AdamW( - params=self.model.parameters(), - **optimizer_kwargs, - ) - elif self.train_config.optimizer.lower() == "sgd": # type: ignore - optimizer_kwargs = get_kwargs(optim.SGD) - optimizer_kwargs.update({k: self.train_config.get(k, v) for k, v in optimizer_kwargs.items()}) - optimizer_kwargs.update(dict(lr=self.lr)) - self.optimizer = optim.SGD( + self.optimizer = opt_cls( params=self.model.parameters(), **optimizer_kwargs, ) else: raise NotImplementedError( - f"optimizer `{self.train_config.optimizer}` not implemented! " # type: ignore - "Please use one of the following: `adam`, `adamw`, `adamw_amsgrad`, `sgd`, " + f"optimizer `{self.train_config.optimizer}` not implemented! " + "Please use one of the optimizers in `torch.optim`, " "or override this method to setup your own optimizer." ) def _setup_scheduler(self) -> None: """Setup the learning rate scheduler.""" - if self.train_config.lr_scheduler is None or self.train_config.lr_scheduler.lower() == "none": # type: ignore + if self.train_config.lr_scheduler is None or self.train_config.lr_scheduler.lower() == "none": self.train_config.lr_scheduler = "none" self.scheduler = None - elif self.train_config.lr_scheduler.lower() == "plateau": # type: ignore + return + + lrs_name = self.train_config.lr_scheduler.lower() + if lrs_name == "onecycle": + lrs_name = "one_cycle" + + if lrs_name not in SCHEDULERS: + # try to find in torch.optim.lr_scheduler + for name in dir(optim.lr_scheduler): + if name.lower() == lrs_name.replace("_", ""): + SCHEDULERS.register(name=lrs_name)(getattr(optim.lr_scheduler, name)) + break + # special case for OneCycleLR + if lrs_name == "one_cycle" and "OneCycleLR" in dir(optim.lr_scheduler): + SCHEDULERS.register(name="one_cycle")(optim.lr_scheduler.OneCycleLR) + + if lrs_name == "plateau": self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.optimizer, "max", patience=2, ) - elif self.train_config.lr_scheduler.lower() == "step": # type: ignore + elif lrs_name == "step": self.scheduler = optim.lr_scheduler.StepLR( self.optimizer, - self.train_config.lr_step_size, # type: ignore - self.train_config.lr_gamma, # type: ignore + self.train_config.lr_step_size, + self.train_config.lr_gamma, ) - elif self.train_config.lr_scheduler.lower() in [ # type: ignore - "one_cycle", - "onecycle", - ]: + elif lrs_name == "one_cycle": self.scheduler = optim.lr_scheduler.OneCycleLR( optimizer=self.optimizer, - max_lr=self.train_config.max_lr, # type: ignore + max_lr=self.train_config.max_lr, epochs=self.n_epochs, - steps_per_epoch=len(self.train_loader), # type: ignore - # verbose=False, + steps_per_epoch=len(self.train_loader), ) - else: # TODO: add linear and linear with warmup schedulers + elif lrs_name in SCHEDULERS: + lrs_cls = SCHEDULERS.get(lrs_name) + lrs_kwargs = get_kwargs(lrs_cls) + lrs_kwargs.update({k: self.train_config.get(k, v) for k, v in lrs_kwargs.items()}) + self.scheduler = lrs_cls(self.optimizer, **lrs_kwargs) + else: raise NotImplementedError( - f"lr scheduler `{self.train_config.lr_scheduler.lower()}` not implemented for training! " # type: ignore - "Please use one of the following: `none`, `plateau`, `step`, `one_cycle`, " + f"lr scheduler `{self.train_config.lr_scheduler}` not implemented for training! " + "Please use one of the schedulers in `torch.optim.lr_scheduler`, " "or override this method to setup your own lr scheduler." ) diff --git a/torch_ecg/models/__init__.py b/torch_ecg/models/__init__.py index 9be54fa7..50b401fd 100644 --- a/torch_ecg/models/__init__.py +++ b/torch_ecg/models/__init__.py @@ -54,11 +54,15 @@ from .ecg_crnn import ECG_CRNN from .ecg_seq_lab_net import ECG_SEQ_LAB_NET from .grad_cam import GradCam +from .registry import BACKBONES, MODELS from .rr_lstm import RR_LSTM from .transformers import Transformer from .unets import ECG_SUBTRACT_UNET, ECG_UNET __all__ = [ + # registries + "BACKBONES", + "MODELS", # CNN backbone "ResNet", "RegNet", diff --git a/torch_ecg/models/cnn/densenet.py b/torch_ecg/models/cnn/densenet.py index b774602c..5fa89995 100644 --- a/torch_ecg/models/cnn/densenet.py +++ b/torch_ecg/models/cnn/densenet.py @@ -25,6 +25,7 @@ from ...models._nets import Conv_Bn_Activation, DownSample from ...utils.misc import CitationMixin, add_docstring, list_sum from ...utils.utils_nn import SizeMixin, compute_sequential_output_shape, compute_sequential_output_shape_docstring +from ..registry import BACKBONES __all__ = [ "DenseNet", @@ -566,6 +567,9 @@ def compute_output_shape( return compute_sequential_output_shape(self, seq_len, batch_size) +@BACKBONES.register(name="densenet") +@BACKBONES.register(name="dense_net") +@BACKBONES.register() class DenseNet(nn.Sequential, SizeMixin, CitationMixin): """The core part of the SOTA model (framework) of CPSC2020. diff --git a/torch_ecg/models/cnn/mobilenet.py b/torch_ecg/models/cnn/mobilenet.py index 86fee0a9..cd6e476b 100644 --- a/torch_ecg/models/cnn/mobilenet.py +++ b/torch_ecg/models/cnn/mobilenet.py @@ -22,6 +22,7 @@ from ...models._nets import Conv_Bn_Activation, Initializers, MultiConv, make_attention_layer from ...utils.misc import CitationMixin, add_docstring from ...utils.utils_nn import SizeMixin, compute_sequential_output_shape, compute_sequential_output_shape_docstring +from ..registry import BACKBONES __all__ = [ "MobileNetV1", @@ -192,6 +193,9 @@ def compute_output_shape( return compute_sequential_output_shape(self, seq_len, batch_size) +@BACKBONES.register(name="mobilenet_v1") +@BACKBONES.register(name="mobile_net_v1") +@BACKBONES.register() class MobileNetV1(nn.Sequential, SizeMixin, CitationMixin): """MobileNet V1. @@ -639,6 +643,9 @@ def compute_output_shape( return output_shape +@BACKBONES.register(name="mobilenet_v2") +@BACKBONES.register(name="mobile_net_v2") +@BACKBONES.register() class MobileNetV2(nn.Sequential, SizeMixin, CitationMixin): """MobileNet V2. @@ -1070,6 +1077,9 @@ def compute_output_shape( return compute_sequential_output_shape(self, seq_len, batch_size) +@BACKBONES.register(name="mobilenet_v3") +@BACKBONES.register(name="mobile_net_v3") +@BACKBONES.register() class MobileNetV3(nn.Sequential, SizeMixin, CitationMixin): """MobileNet V3. diff --git a/torch_ecg/models/cnn/multi_scopic.py b/torch_ecg/models/cnn/multi_scopic.py index fb5b61bf..e999a937 100644 --- a/torch_ecg/models/cnn/multi_scopic.py +++ b/torch_ecg/models/cnn/multi_scopic.py @@ -17,6 +17,7 @@ from ...models._nets import Conv_Bn_Activation, DownSample from ...utils.misc import CitationMixin, add_docstring, list_sum from ...utils.utils_nn import SizeMixin, compute_sequential_output_shape, compute_sequential_output_shape_docstring +from ..registry import BACKBONES __all__ = [ "MultiScopicCNN", @@ -351,6 +352,9 @@ def in_channels(self) -> int: return self.__in_channels +@BACKBONES.register(name="multi_scopic") +@BACKBONES.register(name="multi_scopic_cnn") +@BACKBONES.register() class MultiScopicCNN(nn.Module, SizeMixin, CitationMixin): """CNN part of the SOTA model from CPSC2019 challenge (entry 0416). diff --git a/torch_ecg/models/cnn/regnet.py b/torch_ecg/models/cnn/regnet.py index b0e3795b..0f3cab31 100644 --- a/torch_ecg/models/cnn/regnet.py +++ b/torch_ecg/models/cnn/regnet.py @@ -17,6 +17,7 @@ from ...models._nets import Conv_Bn_Activation, DownSample, SpaceToDepth from ...utils.misc import CitationMixin, add_docstring from ...utils.utils_nn import SizeMixin, compute_sequential_output_shape, compute_sequential_output_shape_docstring +from ..registry import BACKBONES from .resnet import ResNetBasicBlock, ResNetBottleNeck @@ -263,6 +264,8 @@ def compute_output_shape( return compute_sequential_output_shape(self, seq_len, batch_size) +@BACKBONES.register(name="regnet") +@BACKBONES.register() class RegNet(nn.Sequential, SizeMixin, CitationMixin): """RegNet model. diff --git a/torch_ecg/models/cnn/resnet.py b/torch_ecg/models/cnn/resnet.py index 39122731..314ed107 100644 --- a/torch_ecg/models/cnn/resnet.py +++ b/torch_ecg/models/cnn/resnet.py @@ -16,6 +16,7 @@ from ...models._nets import Activations, Conv_Bn_Activation, DownSample, SpaceToDepth, ZeroPadding, make_attention_layer from ...utils.misc import CitationMixin, add_docstring from ...utils.utils_nn import SizeMixin, compute_sequential_output_shape, compute_sequential_output_shape_docstring +from ..registry import BACKBONES __all__ = [ "ResNet", @@ -625,6 +626,9 @@ def compute_output_shape( return compute_sequential_output_shape(self, seq_len, batch_size) +@BACKBONES.register(name="resnet") +@BACKBONES.register(name="resnext") +@BACKBONES.register() class ResNet(nn.Sequential, SizeMixin, CitationMixin): """ResNet model. diff --git a/torch_ecg/models/cnn/vgg.py b/torch_ecg/models/cnn/vgg.py index cad9e128..9157ef15 100644 --- a/torch_ecg/models/cnn/vgg.py +++ b/torch_ecg/models/cnn/vgg.py @@ -15,6 +15,7 @@ compute_sequential_output_shape, compute_sequential_output_shape_docstring, ) +from ..registry import BACKBONES __all__ = [ "VGG16", @@ -127,6 +128,8 @@ def compute_output_shape( return output_shape +@BACKBONES.register(name="vgg16") +@BACKBONES.register() class VGG16(nn.Sequential, SizeMixin, CitationMixin): """CNN feature extractor of VGG architecture. diff --git a/torch_ecg/models/cnn/xception.py b/torch_ecg/models/cnn/xception.py index 0900a5ef..e49e1dab 100644 --- a/torch_ecg/models/cnn/xception.py +++ b/torch_ecg/models/cnn/xception.py @@ -16,6 +16,7 @@ from ...models._nets import DownSample, MultiConv from ...utils.misc import CitationMixin, add_docstring from ...utils.utils_nn import SizeMixin, compute_sequential_output_shape, compute_sequential_output_shape_docstring +from ..registry import BACKBONES __all__ = [ "Xception", @@ -805,6 +806,8 @@ def compute_output_shape( return output_shape +@BACKBONES.register(name="xception") +@BACKBONES.register() class Xception(nn.Sequential, SizeMixin, CitationMixin): """Xception model. diff --git a/torch_ecg/models/ecg_crnn.py b/torch_ecg/models/ecg_crnn.py index 8db5245e..78d1a516 100644 --- a/torch_ecg/models/ecg_crnn.py +++ b/torch_ecg/models/ecg_crnn.py @@ -19,20 +19,16 @@ from ..utils.misc import CitationMixin from ..utils.utils_nn import CkptMixin, SizeMixin from ._nets import MLP, GlobalContextBlock, NonLocalBlock, SEBlock, SelfAttention, StackedLSTM -from .cnn.densenet import DenseNet -from .cnn.mobilenet import MobileNetV1, MobileNetV2, MobileNetV3 -from .cnn.multi_scopic import MultiScopicCNN -from .cnn.regnet import RegNet -from .cnn.resnet import ResNet -from .cnn.vgg import VGG16 -from .cnn.xception import Xception +from .registry import BACKBONES, MODELS from .transformers import Transformer __all__ = [ "ECG_CRNN", + "ECG_CRNN_v1", ] +@MODELS.register() class ECG_CRNN(nn.Module, CkptMixin, SizeMixin, CitationMixin): """Convolutional (Recurrent) Neural Network for ECG tasks. @@ -82,29 +78,20 @@ def __init__( cnn_choice = self.config.cnn.name.lower() # type: ignore cnn_config = self.config.cnn[self.config.cnn.name] # type: ignore - if "resnet" in cnn_choice or "resnext" in cnn_choice: - self.cnn = ResNet(self.n_leads, **cnn_config) - elif "regnet" in cnn_choice: - self.cnn = RegNet(self.n_leads, **cnn_config) - elif "multi_scopic" in cnn_choice: - self.cnn = MultiScopicCNN(self.n_leads, **cnn_config) - elif "mobile_net" in cnn_choice or "mobilenet" in cnn_choice: - if "v1" in cnn_choice: - self.cnn = MobileNetV1(self.n_leads, **cnn_config) - elif "v2" in cnn_choice: - self.cnn = MobileNetV2(self.n_leads, **cnn_config) - elif "v3" in cnn_choice: - self.cnn = MobileNetV3(self.n_leads, **cnn_config) - else: - raise ValueError(f"CNN \042{cnn_choice}\042 is not supported for {self.__name__}") - elif "densenet" in cnn_choice or "dense_net" in cnn_choice: - self.cnn = DenseNet(self.n_leads, **cnn_config) - elif "vgg16" in cnn_choice: - self.cnn = VGG16(self.n_leads, **cnn_config) - elif "xception" in cnn_choice: - self.cnn = Xception(self.n_leads, **cnn_config) - else: + + self.cnn = None + # order by length descending to match the most specific name first + for name in sorted(BACKBONES.list_all(), key=len, reverse=True): + if name.lower() in cnn_choice: + try: + self.cnn = BACKBONES.build(name, in_channels=self.n_leads, **cnn_config) + except TypeError: + self.cnn = BACKBONES.build(name, n_leads=self.n_leads, **cnn_config) + break + + if self.cnn is None: raise NotImplementedError(f"CNN \042{cnn_choice}\042 not implemented yet") + rnn_input_size = self.cnn.compute_output_shape(None, None)[1] if self.config.rnn.name.lower() == "none": # type: ignore @@ -518,29 +505,20 @@ def __init__( cnn_choice = self.config.cnn.name.lower() # type: ignore cnn_config = self.config.cnn[self.config.cnn.name] # type: ignore - if "resnet" in cnn_choice or "resnext" in cnn_choice: - self.cnn = ResNet(self.n_leads, **cnn_config) - elif "regnet" in cnn_choice: - self.cnn = RegNet(self.n_leads, **cnn_config) - elif "multi_scopic" in cnn_choice: - self.cnn = MultiScopicCNN(self.n_leads, **cnn_config) - elif "mobile_net" in cnn_choice or "mobilenet" in cnn_choice: - if "v1" in cnn_choice: - self.cnn = MobileNetV1(self.n_leads, **cnn_config) - elif "v2" in cnn_choice: - self.cnn = MobileNetV2(self.n_leads, **cnn_config) - elif "v3" in cnn_choice: - self.cnn = MobileNetV3(self.n_leads, **cnn_config) - else: - raise ValueError(f"CNN \042{cnn_choice}\042 is not supported for {self.__name__}") - elif "densenet" in cnn_choice or "dense_net" in cnn_choice: - self.cnn = DenseNet(self.n_leads, **cnn_config) - elif "vgg16" in cnn_choice: - self.cnn = VGG16(self.n_leads, **cnn_config) - elif "xception" in cnn_choice: - self.cnn = Xception(self.n_leads, **cnn_config) - else: + + self.cnn = None + # order by length descending to match the most specific name first + for name in sorted(BACKBONES.list_all(), key=len, reverse=True): + if name.lower() in cnn_choice: + try: + self.cnn = BACKBONES.build(name, in_channels=self.n_leads, **cnn_config) + except TypeError: + self.cnn = BACKBONES.build(name, n_leads=self.n_leads, **cnn_config) + break + + if self.cnn is None: raise NotImplementedError(f"CNN \042{cnn_choice}\042 not implemented yet") + rnn_input_size = self.cnn.compute_output_shape(None, None)[1] if self.config.rnn.name.lower() == "none": # type: ignore diff --git a/torch_ecg/models/ecg_seq_lab_net.py b/torch_ecg/models/ecg_seq_lab_net.py index 05a275cf..91147a75 100644 --- a/torch_ecg/models/ecg_seq_lab_net.py +++ b/torch_ecg/models/ecg_seq_lab_net.py @@ -12,12 +12,14 @@ from ..model_configs.ecg_seq_lab_net import ECG_SEQ_LAB_NET_CONFIG from ..utils import add_docstring from .ecg_crnn import ECG_CRNN, ECG_CRNN_v1 +from .registry import MODELS __all__ = [ "ECG_SEQ_LAB_NET", ] +@MODELS.register() class ECG_SEQ_LAB_NET(ECG_CRNN): """SOTA model from CPSC2019 challenge. diff --git a/torch_ecg/models/loss.py b/torch_ecg/models/loss.py index 66ecc341..0c3e5eac 100644 --- a/torch_ecg/models/loss.py +++ b/torch_ecg/models/loss.py @@ -31,6 +31,8 @@ import torch.nn.functional as F from torch import Tensor, nn +from ..components.registry import LOSSES + __all__ = [ "WeightedBCELoss", "BCEWithLogitsWithClassWeightLoss", @@ -99,6 +101,7 @@ def weighted_binary_cross_entropy( return loss.sum() +@LOSSES.register() class WeightedBCELoss(nn.Module): """Weighted Binary Cross Entropy Loss class. @@ -184,6 +187,7 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor: ) +@LOSSES.register() class BCEWithLogitsWithClassWeightLoss(nn.BCEWithLogitsLoss): """Class-weighted Binary Cross Entropy Loss class. @@ -223,6 +227,7 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor: return loss +@LOSSES.register() class MaskedBCEWithLogitsLoss(nn.BCEWithLogitsLoss): """Masked Binary Cross Entropy Loss class. @@ -290,6 +295,7 @@ def forward(self, input: Tensor, target: Tensor, weight_mask: Tensor) -> Tensor: return loss +@LOSSES.register() class FocalLoss(nn.modules.loss._WeightedLoss): """Focal loss class. @@ -409,6 +415,7 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor: return fl +@LOSSES.register() class AsymmetricLoss(nn.Module): """Asymmetric loss class. @@ -646,17 +653,9 @@ def setup_criterion(name: str, **kwargs: Any) -> nn.Module: The criterion (loss function). """ - if name == "WeightedBCELoss": - criterion = WeightedBCELoss(**kwargs) - elif name == "BCEWithLogitsWithClassWeightLoss": - criterion = BCEWithLogitsWithClassWeightLoss(**kwargs) - elif name == "MaskedBCEWithLogitsLoss": - criterion = MaskedBCEWithLogitsLoss(**kwargs) - elif name == "FocalLoss": - criterion = FocalLoss(**kwargs) - elif name == "AsymmetricLoss": - criterion = AsymmetricLoss(**kwargs) - elif name.startswith("nn."): + if name in LOSSES: + return LOSSES.build(name, **kwargs) + if name.startswith("nn."): criterion = eval(name)(**kwargs) elif name in nn.modules.loss.__all__: criterion = getattr(nn, name)(**kwargs) diff --git a/torch_ecg/models/registry.py b/torch_ecg/models/registry.py new file mode 100644 index 00000000..d6164d7c --- /dev/null +++ b/torch_ecg/models/registry.py @@ -0,0 +1,13 @@ +""" +Registries for models and backbones. +""" + +from ..utils.registry import Registry + +__all__ = [ + "BACKBONES", + "MODELS", +] + +BACKBONES = Registry("backbones") +MODELS = Registry("models") diff --git a/torch_ecg/models/rr_lstm.py b/torch_ecg/models/rr_lstm.py index ecd86086..78ca11e7 100644 --- a/torch_ecg/models/rr_lstm.py +++ b/torch_ecg/models/rr_lstm.py @@ -23,12 +23,14 @@ from ..models._nets import MLP, ExtendedCRF, GlobalContextBlock, NonLocalBlock, SEBlock, SelfAttention, StackedLSTM from ..utils.misc import CitationMixin from ..utils.utils_nn import CkptMixin, SizeMixin +from .registry import MODELS __all__ = [ "RR_LSTM", ] +@MODELS.register() class RR_LSTM(nn.Module, CkptMixin, SizeMixin, CitationMixin): """LSTM model for RR time series classification or sequence labeling. diff --git a/torch_ecg/models/unets/ecg_subtract_unet.py b/torch_ecg/models/unets/ecg_subtract_unet.py index f52dad11..fcfb9432 100644 --- a/torch_ecg/models/unets/ecg_subtract_unet.py +++ b/torch_ecg/models/unets/ecg_subtract_unet.py @@ -30,6 +30,7 @@ compute_sequential_output_shape, compute_sequential_output_shape_docstring, ) +from ..registry import MODELS __all__ = [ "ECG_SUBTRACT_UNET", @@ -476,6 +477,7 @@ def compute_output_shape( return output_shape +@MODELS.register() class ECG_SUBTRACT_UNET(nn.Module, CkptMixin, SizeMixin): """U-Net for ECG wave delineation. diff --git a/torch_ecg/models/unets/ecg_unet.py b/torch_ecg/models/unets/ecg_unet.py index 7f048103..2151bf3c 100644 --- a/torch_ecg/models/unets/ecg_unet.py +++ b/torch_ecg/models/unets/ecg_unet.py @@ -29,6 +29,7 @@ compute_sequential_output_shape, compute_sequential_output_shape_docstring, ) +from ..registry import MODELS __all__ = [ "ECG_UNET", @@ -375,6 +376,7 @@ def compute_output_shape( return output_shape +@MODELS.register() class ECG_UNET(nn.Module, CkptMixin, SizeMixin, CitationMixin): """ U-Net for (multi-lead) ECG wave delineation. diff --git a/torch_ecg/preprocessors/__init__.py b/torch_ecg/preprocessors/__init__.py index 9c2b4e25..4f5b917a 100644 --- a/torch_ecg/preprocessors/__init__.py +++ b/torch_ecg/preprocessors/__init__.py @@ -15,6 +15,7 @@ :toctree: generated/ :recursive: + PREPROCESSORS PreprocManager BandPass BaselineRemove @@ -30,9 +31,11 @@ from .baseline_remove import BaselineRemove from .normalize import MinMaxNormalize, NaiveNormalize, Normalize, ZScoreNormalize from .preproc_manager import PreprocManager +from .registry import PREPROCESSORS from .resample import Resample __all__ = [ + "PREPROCESSORS", "PreprocManager", "BandPass", "BaselineRemove", diff --git a/torch_ecg/preprocessors/bandpass.py b/torch_ecg/preprocessors/bandpass.py index 461f8d3c..e73db244 100644 --- a/torch_ecg/preprocessors/bandpass.py +++ b/torch_ecg/preprocessors/bandpass.py @@ -6,12 +6,15 @@ import torch from .._preprocessors.base import preprocess_multi_lead_signal +from .registry import PREPROCESSORS __all__ = [ "BandPass", ] +@PREPROCESSORS.register(name="bandpass") +@PREPROCESSORS.register() class BandPass(torch.nn.Module): """Bandpass filtering preprocessor. diff --git a/torch_ecg/preprocessors/baseline_remove.py b/torch_ecg/preprocessors/baseline_remove.py index a64827ff..49d878ae 100644 --- a/torch_ecg/preprocessors/baseline_remove.py +++ b/torch_ecg/preprocessors/baseline_remove.py @@ -7,12 +7,15 @@ import torch from .._preprocessors.base import preprocess_multi_lead_signal +from .registry import PREPROCESSORS __all__ = [ "BaselineRemove", ] +@PREPROCESSORS.register(name="baseline_remove") +@PREPROCESSORS.register() class BaselineRemove(torch.nn.Module): """Baseline removal using median filtering. diff --git a/torch_ecg/preprocessors/normalize.py b/torch_ecg/preprocessors/normalize.py index faec75ce..dff1a455 100644 --- a/torch_ecg/preprocessors/normalize.py +++ b/torch_ecg/preprocessors/normalize.py @@ -6,6 +6,7 @@ import torch from ..utils.utils_signal_t import normalize as normalize_t +from .registry import PREPROCESSORS __all__ = [ "Normalize", @@ -15,6 +16,8 @@ ] +@PREPROCESSORS.register(name="normalize") +@PREPROCESSORS.register() class Normalize(torch.nn.Module): """Normalization preprocessor. @@ -107,6 +110,8 @@ def forward(self, sig: torch.Tensor) -> torch.Tensor: return sig +@PREPROCESSORS.register(name="min_max_normalize") +@PREPROCESSORS.register() class MinMaxNormalize(Normalize): """Min-Max normalization. @@ -131,6 +136,8 @@ def __init__(self, per_channel: bool = False, inplace: bool = True, **kwargs: An super().__init__(method="min-max", per_channel=per_channel, inplace=inplace, **kwargs) +@PREPROCESSORS.register(name="naive_normalize") +@PREPROCESSORS.register() class NaiveNormalize(Normalize): """Naive normalization @@ -172,6 +179,8 @@ def __init__( ) +@PREPROCESSORS.register(name="z_score_normalize") +@PREPROCESSORS.register() class ZScoreNormalize(Normalize): """Z-score normalization. diff --git a/torch_ecg/preprocessors/preproc_manager.py b/torch_ecg/preprocessors/preproc_manager.py index 93d89182..5a2a7686 100644 --- a/torch_ecg/preprocessors/preproc_manager.py +++ b/torch_ecg/preprocessors/preproc_manager.py @@ -2,16 +2,13 @@ import warnings from random import sample -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple, Union import torch import torch.nn as nn from ..utils.misc import ReprMixin -from .bandpass import BandPass -from .baseline_remove import BaselineRemove -from .normalize import Normalize -from .resample import Resample +from .registry import PREPROCESSORS __all__ = [ "PreprocManager", @@ -59,67 +56,34 @@ def __init__( ) -> None: super().__init__() self.random = random - self._preprocessors = list(pps) + self._preprocessors = nn.ModuleList(list(pps)) - def _add_bandpass(self, **config: dict) -> None: - """Add a bandpass filter to the manager. + def add_(self, pp: Union[nn.Module, str, dict], pos: int = -1, **kwargs: Any) -> None: + """Add a preprocessor to the manager. Parameters ---------- - config : dict - The configuration of the bandpass filter. - - Returns - ------- - None - - """ - self._preprocessors.append(BandPass(**config)) - - def _add_baseline_remove(self, **config: dict) -> None: - """Add a median filter for baseline removal to the manager. - - Parameters - ---------- - config : dict - The configuration of the median filter. - - Returns - ------- - None - - """ - self._preprocessors.append(BaselineRemove(**config)) - - def _add_normalize(self, **config: dict) -> None: - """Add a normalizer to the manager. - - Parameters - ---------- - config : dict - The configuration of the normalizer. - - Returns - ------- - None - - """ - self._preprocessors.append(Normalize(**config)) - - def _add_resample(self, **config: dict) -> None: - """Add a resampler to the manager. - - Parameters - ---------- - config : dict - The configuration of the resampler. - - Returns - ------- - None + pp : torch.nn.Module or str or dict + The preprocessor to be added. + If it's a string or a dict, it will be built using the registry. + pos : int, default -1 + The position to insert the preprocessor. + Should be >= -1, with -1 being the indicator of the end. + **kwargs : Any + Additional keyword arguments for building the preprocessor. """ - self._preprocessors.append(Resample(**config)) + if isinstance(pp, (str, dict)): + pp = PREPROCESSORS.build(pp, **kwargs) + assert isinstance(pp, nn.Module) + assert pp.__class__.__name__ not in [ + p.__class__.__name__ for p in self.preprocessors + ], f"Preprocessor {pp.__class__.__name__} already exists." + assert isinstance(pos, int) and pos >= -1, f"pos must be an integer >= -1, but got {pos}." + if pos == -1: + self._preprocessors.append(pp) + else: + self._preprocessors.insert(pos, pp) def forward(self, sig: torch.Tensor) -> torch.Tensor: """Apply the preprocessors to the signal tensor. @@ -163,20 +127,23 @@ def from_config(cls, config: dict) -> "PreprocManager": """ ppm = cls(random=config.get("random", False), inplace=config.get("inplace", True)) - _mapping = { - "bandpass": ppm._add_bandpass, - "baseline_remove": ppm._add_baseline_remove, - "normalize": ppm._add_normalize, - "resample": ppm._add_resample, - } for pp_name, pp_config in config.items(): if pp_name in [ "random", "inplace", + "fs", ]: continue - if pp_name in _mapping and isinstance(pp_config, dict): - _mapping[pp_name](**pp_config) + if pp_name in PREPROCESSORS or pp_name in [ + "".join([w.capitalize() for w in k.split("_")]) for k in PREPROCESSORS.list_all() + ]: + if isinstance(pp_config, dict): + # add default fs from config if not specified in pp_config + if "fs" not in pp_config and "fs" in config: + pp_config["fs"] = config["fs"] + ppm.add_(pp_name, **pp_config) + else: + ppm.add_(pp_name) else: # just ignore the other items pass @@ -206,6 +173,7 @@ def rearrange(self, new_ordering: List[str]) -> None: "The preprocessors are applied in random order, " "rearranging the preprocessors will not take effect.", RuntimeWarning, ) + # TODO: use a more robust way to map names to classes _mapping = { # built-in preprocessors "Resample": "resample", "BandPass": "bandpass", @@ -220,36 +188,12 @@ def rearrange(self, new_ordering: List[str]) -> None: _mapping.update({k: k}) assert len(new_ordering) == len(set(new_ordering)), "Duplicate preprocessor names." assert len(new_ordering) == len(self._preprocessors), "Number of preprocessors mismatch." - self._preprocessors.sort(key=lambda item: new_ordering.index(_mapping[item.__class__.__name__])) - - def add_(self, pp: nn.Module, pos: int = -1) -> None: - """Add a (custom) preprocessor to the manager. - - This method is preferred against directly manipulating - the internal list of preprocessors via - ``PreprocManager.preprocessors.append(pp)``. - - Parameters - ---------- - pp : torch.nn.Module - The preprocessor to be added. - pos : int, default -1 - The position to insert the preprocessor. - Should be >= -1, with -1 being the indicator of the end. - - """ - assert isinstance(pp, nn.Module) - assert pp.__class__.__name__ not in [ - p.__class__.__name__ for p in self.preprocessors - ], f"Preprocessor {pp.__class__.__name__} already exists." - assert isinstance(pos, int) and pos >= -1, f"pos must be an integer >= -1, but got {pos}." - if pos == -1: - self._preprocessors.append(pp) - else: - self._preprocessors.insert(pos, pp) + pps = list(self._preprocessors) + pps.sort(key=lambda item: new_ordering.index(_mapping[item.__class__.__name__])) + self._preprocessors = nn.ModuleList(pps) @property - def preprocessors(self) -> List[nn.Module]: + def preprocessors(self) -> nn.ModuleList: return self._preprocessors @property diff --git a/torch_ecg/preprocessors/registry.py b/torch_ecg/preprocessors/registry.py new file mode 100644 index 00000000..1e7e8e9d --- /dev/null +++ b/torch_ecg/preprocessors/registry.py @@ -0,0 +1,11 @@ +""" +Registry for preprocessors. +""" + +from ..utils.registry import Registry + +__all__ = [ + "PREPROCESSORS", +] + +PREPROCESSORS = Registry("preprocessors") diff --git a/torch_ecg/preprocessors/resample.py b/torch_ecg/preprocessors/resample.py index b90da8f1..52cd33a1 100644 --- a/torch_ecg/preprocessors/resample.py +++ b/torch_ecg/preprocessors/resample.py @@ -5,12 +5,15 @@ import torch from ..utils.utils_signal_t import resample as resample_t +from .registry import PREPROCESSORS __all__ = [ "Resample", ] +@PREPROCESSORS.register(name="resample") +@PREPROCESSORS.register() class Resample(torch.nn.Module): """Resample the signal into fixed sampling frequency or length. diff --git a/torch_ecg/utils/__init__.py b/torch_ecg/utils/__init__.py index b8a4adb5..f77afdb6 100644 --- a/torch_ecg/utils/__init__.py +++ b/torch_ecg/utils/__init__.py @@ -156,6 +156,7 @@ timeout np_topk select_k + Registry """ @@ -186,6 +187,7 @@ str2bool, timeout, ) +from .registry import Registry from .utils_data import ( ECGWaveForm, ECGWaveFormNames, @@ -277,6 +279,7 @@ "make_serializable", "np_topk", "select_k", + "Registry", "get_mask", "class_weight_to_sample_weight", "ensure_lead_fmt", diff --git a/torch_ecg/utils/registry.py b/torch_ecg/utils/registry.py new file mode 100644 index 00000000..0aa68b8e --- /dev/null +++ b/torch_ecg/utils/registry.py @@ -0,0 +1,137 @@ +""" +Generic registry implementation for managing and building modules (backbones, models, etc.) +""" + +from typing import Any, Dict, List, Optional, Union + +__all__ = [ + "Registry", +] + + +class Registry: + """Registry for managing and building modules. + + A registry is used to map strings (module names) to classes, and provides + a unified interface to instantiate modules from configurations. + + Parameters + ---------- + name : str + Name of the registry. + + Examples + -------- + >>> BACKBONES = Registry("backbones") + >>> @BACKBONES.register() + ... class ResNet(nn.Module): + ... def __init__(self, depth): + ... self.depth = depth + >>> # Build from string + >>> model = BACKBONES.build("ResNet", depth=50) + >>> # Build from config dict + >>> model = BACKBONES.build({"name": "ResNet", "depth": 101}) + + """ + + def __init__(self, name: str) -> None: + self._name = name + self._module_dict: Dict[str, type] = {} + + def __len__(self) -> int: + return len(self._module_dict) + + def __contains__(self, name: str) -> bool: + return name in self._module_dict + + def __repr__(self) -> str: + return f"Registry(name={self._name}, items={list(self._module_dict.keys())})" + + @property + def name(self) -> str: + return self._name + + def register(self, name: Optional[str] = None, override: bool = False) -> Any: + """Decorator to register a module. + + Parameters + ---------- + name : str, optional + Name of the module. If not specified, the class name will be used. + override : bool, default False + Whether to override the existing module with the same name. + + """ + + def _register(cls: type) -> type: + _name = name or cls.__name__ + if not override and _name in self._module_dict: + raise KeyError(f"{_name} is already registered in {self._name} registry") + self._module_dict[_name] = cls + return cls + + return _register + + def get(self, name: str) -> Optional[type]: + """Get the module class by name. + + Parameters + ---------- + name : str + Name of the module. + + Returns + ------- + type or None + The registered module class. + + """ + return self._module_dict.get(name) + + def list_all(self) -> List[str]: + """List all registered modules. + + Returns + ------- + List[str] + A list of all registered module names. + + """ + return list(self._module_dict.keys()) + + def build(self, config: Union[str, Dict[str, Any]], **kwargs: Any) -> Any: + """Build a module from a configuration. + + Parameters + ---------- + config : str or dict + Configuration of the module. + If it's a string, it should be the name of the registered module. + If it's a dict, it must contain a "name" key. + **kwargs : Any + Additional keyword arguments passed to the module's constructor. + + Returns + ------- + Any + The instantiated module. + + """ + if isinstance(config, str): + obj_type = config + obj_config = {} + elif isinstance(config, dict): + obj_config = config.copy() + if "name" not in obj_config: + raise KeyError(f"Configuration for {self._name} must contain a 'name' key") + obj_type = obj_config.pop("name") + else: + raise TypeError(f"Config must be a str or dict, but got {type(config)}") + + obj_cls = self.get(obj_type) + if obj_cls is None: + raise KeyError(f"{obj_type} is not registered in the {self._name} registry") + + # Merge config from dict and extra kwargs + final_kwargs = {**obj_config, **kwargs} + return obj_cls(**final_kwargs) From 1f82422c2b2f49e542020cea9e5f2e81e006b65c Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Fri, 13 Mar 2026 13:35:56 +0800 Subject: [PATCH 2/4] enhance robustness and backward compatibility of the save method of CkptMixin --- torch_ecg/utils/utils_nn.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/torch_ecg/utils/utils_nn.py b/torch_ecg/utils/utils_nn.py index 18c3c429..63f84e53 100644 --- a/torch_ecg/utils/utils_nn.py +++ b/torch_ecg/utils/utils_nn.py @@ -1364,13 +1364,14 @@ def save( path.parent.mkdir(parents=True) extra_items = extra_items or {} + excluded_suffixes = (".pt", ".pth", ".ckpt", ".ckpt.pt", ".ckpt.pth", ".pth.tar", ".pt.tar") if use_safetensors is None: - use_safetensors = path.suffix not in [".pth", ".pt"] + use_safetensors = path.name.endswith(excluded_suffixes) is False _model_config = make_safe_globals(self.config) # type: ignore _train_config = make_safe_globals(train_config) - if path.suffix in [".pth", ".pt"]: + if path.name.endswith(excluded_suffixes): if use_safetensors: path = path.with_suffix(".safetensors") warnings.warn( @@ -1404,7 +1405,7 @@ def save( return if use_safetensors: # not single file - # save the model with safetensors into a zip file with the same name as `path` + # save the model with safetensors with the same name as `path` # `model_config` and `train_config` are saved as json files path = path.with_suffix("") path.mkdir(exist_ok=True) From 1dccf7397e70853f66f24fea24b8514a2962b837 Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Fri, 13 Mar 2026 13:36:46 +0800 Subject: [PATCH 3/4] update the registries for the models --- torch_ecg/models/_nets.py | 18 ++ torch_ecg/models/ecg_crnn.py | 235 ++++++++----------- torch_ecg/models/registry.py | 2 + torch_ecg/models/transformers/transformer.py | 4 + 4 files changed, 118 insertions(+), 141 deletions(-) diff --git a/torch_ecg/models/_nets.py b/torch_ecg/models/_nets.py index 4a92825e..66486fe2 100644 --- a/torch_ecg/models/_nets.py +++ b/torch_ecg/models/_nets.py @@ -27,6 +27,7 @@ compute_maxpool_output_shape, compute_receptive_field, ) +from .registry import ATTN_LAYERS __all__ = [ "Initializers", @@ -2457,6 +2458,9 @@ def extra_repr(self) -> str: ) +@ATTN_LAYERS.register(name="self_attention") +@ATTN_LAYERS.register(name="sa") +@ATTN_LAYERS.register() class SelfAttention(nn.Module, SizeMixin): """Self attention layer. @@ -2561,6 +2565,8 @@ def compute_output_shape( return output_shape +@ATTN_LAYERS.register(name="attentive_pooling") +@ATTN_LAYERS.register() class AttentivePooling(nn.Module, SizeMixin): """Attentive pooling layer. @@ -2919,6 +2925,9 @@ def __init__( ) +@ATTN_LAYERS.register(name="non_local") +@ATTN_LAYERS.register(name="nl") +@ATTN_LAYERS.register() class NonLocalBlock(nn.Module, SizeMixin): """Non-local Attention Block. @@ -3045,6 +3054,9 @@ def in_channels(self) -> int: return self.__in_channels +@ATTN_LAYERS.register(name="se_block") +@ATTN_LAYERS.register(name="se") +@ATTN_LAYERS.register() class SEBlock(nn.Module, SizeMixin): """Squeeze-and-Excitation Block. @@ -3183,6 +3195,9 @@ def __init__(self, in_channels: int, **kwargs: Any) -> None: raise NotImplementedError +@ATTN_LAYERS.register(name="global_context") +@ATTN_LAYERS.register(name="gc") +@ATTN_LAYERS.register() class GlobalContextBlock(nn.Module, SizeMixin): """Global Context Block. @@ -3345,6 +3360,9 @@ def __init__(self, in_channels: int, **kwargs: Any): raise NotImplementedError +@ATTN_LAYERS.register(name="cbam_block") +@ATTN_LAYERS.register(name="cbam") +@ATTN_LAYERS.register() class CBAMBlock(nn.Module, SizeMixin): """Convolutional Block Attention Module (ECCV2018). diff --git a/torch_ecg/models/ecg_crnn.py b/torch_ecg/models/ecg_crnn.py index 78d1a516..0d50edf1 100644 --- a/torch_ecg/models/ecg_crnn.py +++ b/torch_ecg/models/ecg_crnn.py @@ -18,9 +18,8 @@ from ..model_configs.ecg_crnn import ECG_CRNN_CONFIG from ..utils.misc import CitationMixin from ..utils.utils_nn import CkptMixin, SizeMixin -from ._nets import MLP, GlobalContextBlock, NonLocalBlock, SEBlock, SelfAttention, StackedLSTM -from .registry import BACKBONES, MODELS -from .transformers import Transformer +from ._nets import MLP, StackedLSTM +from .registry import ATTN_LAYERS, BACKBONES, MODELS __all__ = [ "ECG_CRNN", @@ -92,7 +91,7 @@ def __init__( if self.cnn is None: raise NotImplementedError(f"CNN \042{cnn_choice}\042 not implemented yet") - rnn_input_size = self.cnn.compute_output_shape(None, None)[1] + rnn_input_size = self.cnn.compute_output_shape(2000, 2)[1] if self.config.rnn.name.lower() == "none": # type: ignore self.rnn_in_rearrange = Rearrange("batch_size channels seq_len -> seq_len batch_size channels") @@ -147,77 +146,43 @@ def __init__( self.__attn_seqlen_dim = -1 self.attn_out_rearrange = nn.Identity() clf_input_size = attn_input_size - elif self.config.attn.name.lower() == "nl": # type: ignore - # non_local - self.attn_in_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") - self.attn = NonLocalBlock( - in_channels=attn_input_size, # type: ignore - filter_lengths=self.config.attn.nl.filter_lengths, # type: ignore - subsample_length=self.config.attn.nl.subsample_length, # type: ignore - batch_norm=self.config.attn.nl.batch_norm, # type: ignore - ) - self.__attn_seqlen_dim = -1 - self.attn_out_rearrange = nn.Identity() - clf_input_size = self.attn.compute_output_shape(None, None)[1] - elif self.config.attn.name.lower() == "se": # type: ignore - # squeeze_exitation - self.attn_in_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") - self.attn = SEBlock( - in_channels=attn_input_size, # type: ignore - reduction=self.config.attn.se.reduction, # type: ignore - activation=self.config.attn.se.activation, # type: ignore - kw_activation=self.config.attn.se.kw_activation, # type: ignore - bias=self.config.attn.se.bias, # type: ignore - ) - self.__attn_seqlen_dim = -1 - self.attn_out_rearrange = nn.Identity() - clf_input_size = self.attn.compute_output_shape(None, None)[1] - elif self.config.attn.name.lower() == "gc": # type: ignore - # global_context - self.attn_in_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") - self.attn = GlobalContextBlock( - in_channels=attn_input_size, # type: ignore - ratio=self.config.attn.gc.ratio, # type: ignore - reduction=self.config.attn.gc.reduction, # type: ignore - pooling_type=self.config.attn.gc.pooling_type, # type: ignore - fusion_types=self.config.attn.gc.fusion_types, # type: ignore - ) - self.__attn_seqlen_dim = -1 - self.attn_out_rearrange = nn.Identity() - clf_input_size = self.attn.compute_output_shape(None, None)[1] - elif self.config.attn.name.lower() == "sa": # type: ignore - # self_attention - # NOTE: this branch NOT tested - self.attn_in_rearrange = nn.Identity() - self.attn = SelfAttention( # type: ignore - embed_dim=attn_input_size, - num_heads=self.config.attn.sa.get("num_heads", self.config.attn.sa.get("head_num")), # type: ignore - dropout=self.config.attn.sa.dropout, # type: ignore - bias=self.config.attn.sa.bias, # type: ignore - ) - self.__attn_seqlen_dim = 0 - self.attn_out_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") - clf_input_size = self.attn.compute_output_shape(None, None)[-1] - elif self.config.attn.name.lower() == "transformer": # type: ignore - self.attn = Transformer( - input_size=attn_input_size, # type: ignore - hidden_size=self.config.attn.transformer.hidden_size, # type: ignore - num_layers=self.config.attn.transformer.num_layers, # type: ignore - num_heads=self.config.attn.transformer.num_heads, # type: ignore - dropout=self.config.attn.transformer.dropout, # type: ignore - activation=self.config.attn.transformer.activation, # type: ignore - ) - if self.attn.batch_first: - self.attn_in_rearrange = Rearrange("seq_len batch_size channels -> batch_size seq_len channels") - self.attn_out_rearrange = Rearrange("batch_size seq_len channels -> batch_size channels seq_len") - self.__attn_seqlen_dim = 1 - else: + else: + attn_choice = self.config.attn.name.lower() + attn_config = self.config.attn[self.config.attn.name] + self.attn = None + for name in sorted(ATTN_LAYERS.list_all(), key=len, reverse=True): + if name.lower() in attn_choice: + if name.lower() in ["transformer", "transformer_encoder"]: + self.attn = ATTN_LAYERS.build(name, input_size=attn_input_size, **attn_config) + elif name.lower() in ["sa", "self_attention", "multi_head_attention", "attentive_pooling"]: + self.attn = ATTN_LAYERS.build(name, in_features=attn_input_size, **attn_config) + else: + self.attn = ATTN_LAYERS.build(name, in_channels=attn_input_size, **attn_config) + break + + if self.attn is None: + raise NotImplementedError(f"Attention \042{self.config.attn.name}\042 not implemented yet") + + if attn_choice in ["nl", "non_local", "se", "se_block", "gc", "global_context", "cbam", "cbam_block"]: + self.attn_in_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") + self.__attn_seqlen_dim = -1 + self.attn_out_rearrange = nn.Identity() + clf_input_size = int(self.attn.compute_output_shape(2000, 2)[1]) + elif attn_choice in ["sa", "self_attention"]: self.attn_in_rearrange = nn.Identity() - self.attn_out_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") self.__attn_seqlen_dim = 0 - clf_input_size = self.attn.compute_output_shape(None, None)[-1] - else: - raise NotImplementedError(f"Attention \042{self.config.attn.name}\042 not implemented yet") # type: ignore + self.attn_out_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") + clf_input_size = self.attn.compute_output_shape(2000, 2)[-1] + elif attn_choice in ["transformer", "transformer_encoder"]: + if self.attn.batch_first: + self.attn_in_rearrange = Rearrange("seq_len batch_size channels -> batch_size seq_len channels") + self.attn_out_rearrange = Rearrange("batch_size seq_len channels -> batch_size channels seq_len") + self.__attn_seqlen_dim = 1 + else: + self.attn_in_rearrange = nn.Identity() + self.attn_out_rearrange = Rearrange("seq_len batch_size channels -> batch_size channels seq_len") + self.__attn_seqlen_dim = 0 + clf_input_size = self.attn.compute_output_shape(2000, 2)[-1] # global pooling if self.config.rnn.name.lower() == "lstm" and not self.config.rnn.lstm.retseq: # type: ignore @@ -430,7 +395,14 @@ def from_v1( """ v1_model, train_config = ECG_CRNN_v1.from_checkpoint(v1_ckpt, device=device, weights_only=False) - model = cls(classes=v1_model.classes, n_leads=v1_model.n_leads, config=v1_model.config) + # v1 models usually have no global pooling + # and the classifier input size is the cnn/rnn/attn output size + # which is usually 2000 if seq_len is 2000 + # however, in the new version, the default is max pooling + config = deepcopy(ECG_CRNN_CONFIG) + config.update(deepcopy(v1_model.config)) + config.global_pool = "none" + model = cls(classes=v1_model.classes, n_leads=v1_model.n_leads, config=config) model = model.to(v1_model.device) model.cnn.load_state_dict(v1_model.cnn.state_dict()) if model.rnn.__class__.__name__ != "Identity": @@ -444,6 +416,7 @@ def from_v1( return model +@MODELS.register() class ECG_CRNN_v1(nn.Module, CkptMixin, SizeMixin, CitationMixin): """Convolutional (Recurrent) Neural Network for ECG tasks. @@ -519,7 +492,8 @@ def __init__( if self.cnn is None: raise NotImplementedError(f"CNN \042{cnn_choice}\042 not implemented yet") - rnn_input_size = self.cnn.compute_output_shape(None, None)[1] + rnn_input_size = self.cnn.compute_output_shape(2000, 2)[1] + clf_input_size = rnn_input_size # default if self.config.rnn.name.lower() == "none": # type: ignore self.rnn = None @@ -548,9 +522,10 @@ def __init__( raise NotImplementedError(f"RNN \042{self.config.rnn.name}\042 not implemented yet") # type: ignore # attention + clf_input_size = attn_input_size if self.config.rnn.name.lower() == "lstm" and not self.config.rnn.lstm.retseq: # type: ignore self.attn = None - clf_input_size = attn_input_size + self.__attn_seqlen_dim = 0 if self.config.attn.name.lower() != "none": # type: ignore warnings.warn( f"since `retseq` of rnn is False, hence attention `{self.config.attn.name}` is ignored", # type: ignore @@ -558,58 +533,33 @@ def __init__( ) elif self.config.attn.name.lower() == "none": # type: ignore self.attn = None - clf_input_size = attn_input_size - elif self.config.attn.name.lower() == "nl": # type: ignore - # non_local - self.attn = NonLocalBlock( - in_channels=attn_input_size, # type: ignore - filter_lengths=self.config.attn.nl.filter_lengths, # type: ignore - subsample_length=self.config.attn.nl.subsample_length, # type: ignore - batch_norm=self.config.attn.nl.batch_norm, # type: ignore - ) - clf_input_size = self.attn.compute_output_shape(None, None)[1] - elif self.config.attn.name.lower() == "se": # type: ignore - # squeeze_exitation - self.attn = SEBlock( - in_channels=attn_input_size, # type: ignore - reduction=self.config.attn.se.reduction, # type: ignore - activation=self.config.attn.se.activation, # type: ignore - kw_activation=self.config.attn.se.kw_activation, # type: ignore - bias=self.config.attn.se.bias, # type: ignore - ) - clf_input_size = self.attn.compute_output_shape(None, None)[1] - elif self.config.attn.name.lower() == "gc": # type: ignore - # global_context - self.attn = GlobalContextBlock( - in_channels=attn_input_size, # type: ignore - ratio=self.config.attn.gc.ratio, # type: ignore - reduction=self.config.attn.gc.reduction, # type: ignore - pooling_type=self.config.attn.gc.pooling_type, # type: ignore - fusion_types=self.config.attn.gc.fusion_types, # type: ignore - ) - clf_input_size = self.attn.compute_output_shape(None, None)[1] - elif self.config.attn.name.lower() == "sa": # type: ignore - # self_attention - # NOTE: this branch NOT tested - self.attn = SelfAttention( # type: ignore - embed_dim=attn_input_size, - num_heads=self.config.attn.sa.get("num_heads", self.config.attn.sa.get("head_num")), # type: ignore - dropout=self.config.attn.sa.dropout, # type: ignore - bias=self.config.attn.sa.bias, # type: ignore - ) - clf_input_size = self.attn.compute_output_shape(None, None)[-1] - elif self.config.attn.name.lower() == "transformer": # type: ignore - self.attn = Transformer( - input_size=attn_input_size, # type: ignore - hidden_size=self.config.attn.transformer.hidden_size, # type: ignore - num_layers=self.config.attn.transformer.num_layers, # type: ignore - num_heads=self.config.attn.transformer.num_heads, # type: ignore - dropout=self.config.attn.transformer.dropout, # type: ignore - activation=self.config.attn.transformer.activation, # type: ignore - ) - clf_input_size = self.attn.compute_output_shape(None, None)[-1] + self.__attn_seqlen_dim = -1 else: - raise NotImplementedError(f"Attention \042{self.config.attn.name}\042 not implemented yet") # type: ignore + attn_choice = self.config.attn.name.lower() + attn_config = self.config.attn[self.config.attn.name] + self.attn = None + for name in sorted(ATTN_LAYERS.list_all(), key=len, reverse=True): + if name.lower() in attn_choice: + if name.lower() in ["transformer", "transformer_encoder"]: + self.attn = ATTN_LAYERS.build(name, input_size=attn_input_size, **attn_config) + elif name.lower() in ["sa", "self_attention", "multi_head_attention", "attentive_pooling"]: + self.attn = ATTN_LAYERS.build(name, in_features=attn_input_size, **attn_config) + else: + self.attn = ATTN_LAYERS.build(name, in_channels=attn_input_size, **attn_config) + break + + if self.attn is None: + raise NotImplementedError(f"Attention \042{self.config.attn.name}\042 not implemented yet") + + if attn_choice in ["nl", "non_local", "se", "se_block", "gc", "global_context", "cbam", "cbam_block"]: + clf_input_size = int(self.attn.compute_output_shape(2000, 2)[1]) + self.__attn_seqlen_dim = -1 + elif attn_choice in ["sa", "self_attention"]: + clf_input_size = int(self.attn.compute_output_shape(2000, 2)[-1]) + self.__attn_seqlen_dim = 0 + elif attn_choice in ["transformer", "transformer_encoder"]: + clf_input_size = int(self.attn.compute_output_shape(2000, 2)[-1]) + self.__attn_seqlen_dim = 1 if self.attn.batch_first else 0 if self.config.rnn.name.lower() == "lstm" and not self.config.rnn.lstm.retseq: # type: ignore self.pool = None @@ -689,20 +639,23 @@ def extract_features(self, input: Tensor) -> Tensor: if self.attn is None and features.ndim == 3: # (seq_len, batch_size, channels) --> (batch_size, channels, seq_len) features = features.permute(1, 2, 0) - elif self.config.attn.name.lower() in ["nl", "se", "gc"]: # type: ignore - # (seq_len, batch_size, channels) --> (batch_size, channels, seq_len) - features = features.permute(1, 2, 0) - # (batch_size, channels, seq_len) - features = self.attn(features) # type: ignore - elif self.config.attn.name.lower() in ["sa"]: # type: ignore - # (seq_len, batch_size, channels) - features = self.attn(features) # type: ignore - # (seq_len, batch_size, channels) -> (batch_size, channels, seq_len) - features = features.permute(1, 2, 0) - elif self.config.attn.name.lower() in ["transformer"]: # type: ignore - features = self.attn(features) # type: ignore - # (seq_len, batch_size, channels) -> (batch_size, channels, seq_len) - features = features.permute(1, 2, 0) + elif self.attn is not None: + if self.__attn_seqlen_dim == -1: + # (seq_len, batch_size, channels) --> (batch_size, channels, seq_len) + features = features.permute(1, 2, 0) + # (batch_size, channels, seq_len) + features = self.attn(features) + elif self.__attn_seqlen_dim == 0: + # (seq_len, batch_size, channels) + features = self.attn(features) + # (seq_len, batch_size, channels) -> (batch_size, channels, seq_len) + features = features.permute(1, 2, 0) + elif self.__attn_seqlen_dim == 1: + # (seq_len, batch_size, channels) -> (batch_size, seq_len, channels) + features = features.permute(1, 0, 2) + features = self.attn(features) + # (batch_size, seq_len, channels) -> (batch_size, channels, seq_len) + features = features.permute(0, 2, 1) return features def forward(self, input: Tensor) -> Tensor: diff --git a/torch_ecg/models/registry.py b/torch_ecg/models/registry.py index d6164d7c..5d42564c 100644 --- a/torch_ecg/models/registry.py +++ b/torch_ecg/models/registry.py @@ -7,7 +7,9 @@ __all__ = [ "BACKBONES", "MODELS", + "ATTN_LAYERS", ] BACKBONES = Registry("backbones") MODELS = Registry("models") +ATTN_LAYERS = Registry("attn_layers") diff --git a/torch_ecg/models/transformers/transformer.py b/torch_ecg/models/transformers/transformer.py index 0318f0e7..a8ad5a84 100644 --- a/torch_ecg/models/transformers/transformer.py +++ b/torch_ecg/models/transformers/transformer.py @@ -9,12 +9,16 @@ from ...utils.misc import get_kwargs from ...utils.utils_nn import SizeMixin +from ..registry import ATTN_LAYERS __all__ = [ "Transformer", ] +@ATTN_LAYERS.register(name="transformer") +@ATTN_LAYERS.register(name="transformer_encoder") +@ATTN_LAYERS.register() class Transformer(nn.Module, SizeMixin): """Transformer feature extractor. From a147f346fd5f269b1ad4eebf045956a514f94e84 Mon Sep 17 00:00:00 2001 From: WEN Hao Date: Fri, 13 Mar 2026 22:47:41 +0800 Subject: [PATCH 4/4] fix error in from_config method for augmenter/preprocessor managers --- torch_ecg/augmenters/augmenter_manager.py | 2 ++ torch_ecg/preprocessors/preproc_manager.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/torch_ecg/augmenters/augmenter_manager.py b/torch_ecg/augmenters/augmenter_manager.py index 197522d2..a58177db 100644 --- a/torch_ecg/augmenters/augmenter_manager.py +++ b/torch_ecg/augmenters/augmenter_manager.py @@ -152,6 +152,8 @@ def from_config(cls, config: dict) -> "AugmenterManager": if aug_name in AUGMENTERS or aug_name in [ "".join([w.capitalize() for w in k.split("_")]) for k in AUGMENTERS.list_all() ]: + if aug_config is False: + continue if isinstance(aug_config, dict): # add default fs from config if not specified in aug_config if "fs" not in aug_config and "fs" in config: diff --git a/torch_ecg/preprocessors/preproc_manager.py b/torch_ecg/preprocessors/preproc_manager.py index 5a2a7686..7bdd6bb1 100644 --- a/torch_ecg/preprocessors/preproc_manager.py +++ b/torch_ecg/preprocessors/preproc_manager.py @@ -137,6 +137,8 @@ def from_config(cls, config: dict) -> "PreprocManager": if pp_name in PREPROCESSORS or pp_name in [ "".join([w.capitalize() for w in k.split("_")]) for k in PREPROCESSORS.list_all() ]: + if pp_config is False: + continue if isinstance(pp_config, dict): # add default fs from config if not specified in pp_config if "fs" not in pp_config and "fs" in config: