Implement comprehensive Registry Pattern for system-wide modularity and decoupling#29
Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| 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 |
There was a problem hiding this comment.
Pull request overview
This PR introduces a generic Registry mechanism and refactors multiple subsystems (models/backbones/attention layers, preprocessors, augmenters, trainer components) to construct components dynamically from configuration instead of hardcoded conditionals, aiming to improve modularity and extensibility across torch_ecg.
Changes:
- Added a reusable
Registryutility and introduced registries for preprocessors, augmenters, models/backbones/attention layers, and trainer components (optimizers/schedulers/losses). - Refactored
ECG_CRNN/ECG_CRNN_v1,PreprocManager,AugmenterManager, and trainer optimizer/scheduler/loss setup to use registry-based building. - Registered many existing components via decorators and updated CI/tooling configs.
Reviewed changes
Copilot reviewed 43 out of 43 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| torch_ecg/utils/utils_nn.py | Adjust safetensors vs torch-save selection |
| torch_ecg/utils/registry.py | Add generic Registry implementation |
| torch_ecg/utils/init.py | Export Registry from utils |
| torch_ecg/preprocessors/resample.py | Register Resample in PREPROCESSORS |
| torch_ecg/preprocessors/registry.py | Add PREPROCESSORS registry |
| torch_ecg/preprocessors/preproc_manager.py | Build preprocessors via registry |
| torch_ecg/preprocessors/normalize.py | Register normalization preprocessors |
| torch_ecg/preprocessors/baseline_remove.py | Register BaselineRemove |
| torch_ecg/preprocessors/bandpass.py | Register BandPass |
| torch_ecg/preprocessors/init.py | Export PREPROCESSORS |
| torch_ecg/models/unets/ecg_unet.py | Register ECG_UNET in MODELS |
| torch_ecg/models/unets/ecg_subtract_unet.py | Register ECG_SUBTRACT_UNET in MODELS |
| torch_ecg/models/transformers/transformer.py | Register Transformer in ATTN_LAYERS |
| torch_ecg/models/rr_lstm.py | Register RR_LSTM in MODELS |
| torch_ecg/models/registry.py | Add BACKBONES/MODELS/ATTN_LAYERS registries |
| torch_ecg/models/loss.py | Register losses + use LOSSES in setup |
| torch_ecg/models/ecg_seq_lab_net.py | Register ECG_SEQ_LAB_NET in MODELS |
| torch_ecg/models/ecg_crnn.py | Refactor backbone/attn selection via registries |
| torch_ecg/models/cnn/xception.py | Register Xception backbone |
| torch_ecg/models/cnn/vgg.py | Register VGG16 backbone |
| torch_ecg/models/cnn/resnet.py | Register ResNet/ResNeXt backbones |
| torch_ecg/models/cnn/regnet.py | Register RegNet backbone |
| torch_ecg/models/cnn/multi_scopic.py | Register MultiScopicCNN backbone |
| torch_ecg/models/cnn/mobilenet.py | Register MobileNet backbones |
| torch_ecg/models/cnn/densenet.py | Register DenseNet backbone |
| torch_ecg/models/_nets.py | Register attention blocks in ATTN_LAYERS |
| torch_ecg/models/init.py | Export registries from models package |
| torch_ecg/components/trainer.py | Registry-based optimizer/scheduler resolution |
| torch_ecg/components/registry.py | Add OPTIMIZERS/SCHEDULERS/LOSSES registries |
| torch_ecg/augmenters/stretch_compress.py | Register StretchCompress augmenter |
| torch_ecg/augmenters/registry.py | Add AUGMENTERS registry |
| torch_ecg/augmenters/random_renormalize.py | Register RandomRenormalize augmenter |
| torch_ecg/augmenters/random_masking.py | Register RandomMasking augmenter |
| torch_ecg/augmenters/random_flip.py | Register RandomFlip augmenter |
| torch_ecg/augmenters/mixup.py | Register Mixup augmenter |
| torch_ecg/augmenters/label_smooth.py | Register LabelSmooth augmenter |
| torch_ecg/augmenters/cutmix.py | Register CutMix augmenter |
| torch_ecg/augmenters/baseline_wander.py | Register BaselineWanderAugmenter |
| torch_ecg/augmenters/augmenter_manager.py | Build augmenters via registry |
| torch_ecg/augmenters/init.py | Export AUGMENTERS |
| .pre-commit-config.yaml | Update flake8 ignore list |
| .github/workflows/test-models.yml | Update Python test matrix |
| .github/workflows/test-databases.yml | Update Python test matrix |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| 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 | ||
|
|
| from .random_renormalize import RandomRenormalize | ||
| from .stretch_compress import StretchCompress | ||
| from .registry import AUGMENTERS | ||
|
|
| for name in dir(optim.lr_scheduler): | ||
| if name.lower() == lrs_name.replace("_", ""): |
| if name in LOSSES: | ||
| return LOSSES.build(name, **kwargs) | ||
| if name.startswith("nn."): | ||
| criterion = eval(name)(**kwargs) |
| 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") |
| 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) |
| strategy: | ||
| matrix: | ||
| python-version: ["3.9", "3.10", "3.11", "3.12"] | ||
| python-version: ["3.9", "3.10", "3.12", "3.13"] |
| strategy: | ||
| matrix: | ||
| python-version: ["3.9", "3.10", "3.11", "3.12"] | ||
| python-version: ["3.9", "3.10", "3.12", "3.13"] |
There was a problem hiding this comment.
The test of 3.11 got stuck for an unknown reason. Currently, there's no time to resolve it.
| 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) |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1dccf7397e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| else: | ||
| am.add_(aug_name) |
There was a problem hiding this comment.
AugmenterManager.from_config now instantiates any registered augmenter whose config value is non-dict via am.add_(aug_name), so entries set to False are no longer treated as disabled. Existing training configs commonly use label_smooth = False, random_masking = False, etc. (for example in torch_ecg/databases/datasets/cinc2021/cinc2021_cfg.py), so this change silently enables augmenters with default probabilities and alters training behavior/results.
Useful? React with 👍 / 👎.
| else: | ||
| ppm.add_(pp_name) |
There was a problem hiding this comment.
PreprocManager.from_config has the same non-dict path and now calls ppm.add_(pp_name) even when a preprocessor is explicitly disabled with False. Several dataset configs rely on normalize = False (for example torch_ecg/databases/datasets/cpsc2019/cpsc2019_cfg.py), so this change can unexpectedly add default preprocessors and modify input distributions in experiments that previously skipped them.
Useful? React with 👍 / 👎.
| 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) |
There was a problem hiding this comment.
The generic scheduler path only builds kwargs from get_kwargs(lrs_cls) (defaults) and then updates those same keys, so constructor arguments without defaults can never be passed from train_config. As a result, many torch.optim.lr_scheduler classes that require extra parameters beyond optimizer will fail at runtime with missing-argument errors despite being selected by name.
Useful? React with 👍 / 👎.
| 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) |
There was a problem hiding this comment.
In ECG_CRNN, the registry branch groups attentive_pooling with self-attention and instantiates it using in_features=..., but AttentivePooling expects in_channels. Selecting attn.name = "attentive_pooling" therefore raises an unexpected-keyword TypeError during model initialization.
Useful? React with 👍 / 👎.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #29 +/- ##
==========================================
- Coverage 93.52% 93.46% -0.07%
==========================================
Files 134 139 +5
Lines 18252 18413 +161
==========================================
+ Hits 17071 17210 +139
- Misses 1181 1203 +22 ☔ View full report in Codecov by Sentry. |
This PR introduces a unified Registry system across torch_ecg, eliminating hardcoded if-elif chains in core modules. This architectural upgrade significantly enhances scalability, allowing for dynamic component injection via simple decorators and ensuring a cleaner, more maintainable codebase.
Key Changes
Registryclass intorch_ecg.utils.registryproviding register, get, and build capabilities.BACKBONES,MODELS, andATTN_LAYERSregistries.ECG_CRNNandECG_CRNN_v1to utilize dynamic building for both CNN backbones and Attention layers.@registerdecorators to all major CNN backbones (ResNet,VGG,RegNet,MobileNet, etc.) and downstream models.in_channels,in_features, andembed_dimaliases) to ensure registry compatibility across different attention modules.PREPROCESSORSandAUGMENTERSregistries.PreprocManagerandAugmenterManagerto dynamically build processing pipelines directly from configurations.OPTIMIZERS,SCHEDULERS, andLOSSESregistries.torch.optim,torch.nn).