Merge pull request #38 from DeepPSP/arch-preproc-pytorch

Enhance the preprocessors module with pure pytorch implementations

Author: kjs11

.gitignore

@@ -162,3 +162,5 @@ log
 *.safetensors
 
 tmp*
+
+.github/copilot-instructions.md

CHANGELOG.rst

@@ -34,6 +34,8 @@ Changed
   `torch_ecg.utils.download` module.
 - Add length verification in the `http_get` function in
   `torch_ecg.utils.download` module.
+- Enhance the pytorch `preprocessors` module: all preprocessors now have
+  pure PyTorch implementations, and NumPy fallbacks for NumPy array inputs.
 
 Deprecated
 ~~~~~~~~~~

ROADMAP.md

@@ -4,47 +4,68 @@ This document outlines the architectural upgrades and future development plans f
 
 ## Architectural Upgrades
 
-### 1. Introduction of Registry Pattern
+### 1. Introduction of Registry Pattern ✅ Done (PR #29)
 **Objective**: Eliminate lengthy `if-elif` branches in downstream models like `ECG_CRNN` to enhance code maintainability and extensibility.
-- **Current State**: Adding a new backbone currently requires modifying multiple files, including `models/ecg_crnn.py`.
+- **Status**: Done. `MODELS`, `BACKBONES`, `ATTN_LAYERS` registries are implemented in `models/registry.py`; `OPTIMIZERS`, `SCHEDULERS`, `LOSSES` in `components/registry.py`; `PREPROCESSORS` in `preprocessors/registry.py`. All CNN backbones and downstream models use `@BACKBONES.register()` / `@MODELS.register()` decorators. `Registry.build(name, **kwargs)` is the unified construction interface. Adding a new backbone no longer requires modifying `models/ecg_crnn.py` or any other existing file.
 - **Strategy**: Implement a registry mechanism similar to the one used in the `fl-sim` library. Establish `BACKBONES`, `MODELS`, and `SSL` registries.
   - Use decorators like `@register_backbone("resnet")` to register modules.
   - Use a unified `BACKBONES.build(name, **kwargs)` method for module instantiation.
 - **Benefits**: Decouples model definition from construction logic, making it easier for both maintainers and users to inject custom backbones.
 
-### 2. Standardized Backbone API
+### 2. Standardized Backbone API ✅ Done (PR #30)
 **Objective**: Provide a unified feature extraction interface for Self-Supervised Learning (SSL) and multi-task learning.
+- **Status**: Done. All CNN backbones (`ResNet`, `VGG16`, `DenseNet`, `MobileNetV1/V2/V3`, `MultiScopicCNN`, `RegNet`, `Xception`) and the `Transformer` now implement `forward_features(x)` (returns feature maps before the classifier head) and `compute_features_output_shape(seq_len, batch_size)` for shape inference without running a forward pass.
 - **Strategy**: Follow the convention used in the `timm` library by providing a `forward_features(x)` method for all CNN and Transformer backbones.
 - **Features**:
   - Unified return of feature maps instead of classification logits.
   - Support for accessing intermediate activations for feature fusion or saliency analysis (e.g., Grad-CAM).
 
-### 3. Leveraging Lazy Modules for Configuration Optimization
+### 3. Leveraging Lazy Modules for Configuration Optimization ⬜ Not started
 **Objective**: Reduce the burden of manually calculating and specifying `in_channels` in configuration files.
 - **Strategy**: Introduce `nn.LazyLinear` or `nn.LazyConv1d` in complex SSL modules.
 - **Benefits**: Simplifies `model_configs` by allowing the model to automatically infer input dimensions during the first forward pass, reducing boilerplate code.
 
-### 4. Consolidation and Optimization of Preprocessors and Augmenters
+### 4. Consolidation and Optimization of Preprocessors and Augmenters 🔄 In Progress
 **Objective**: Eliminate redundancy between NumPy and PyTorch implementations and optimize performance by keeping computations on the GPU.
-- **Pure PyTorch Filtering**: Implement `BandPass` and `BaselineRemove` using pure PyTorch (e.g., using `torchaudio.functional` or custom FFT-based filters) to avoid expensive CPU-GPU data transfers.
-- **Unification of Managers**: Refactor `PreprocManager` and `AugmenterManager` to share a common base or registry, as their logic for managing sequences of transforms is very similar.
-- **Dimension Agnostic Transforms**: Ensure all preprocessors and augmenters can handle arbitrary batch and lead dimensions (using `...` in slicing and einops where possible), reducing the need for functions like `preprocess_multi_lead_signal`.
-- **Numpy Version Maintenance**: Keep `_preprocessors` (NumPy version) only for offline data preparation or deployment environments without PyTorch, while making the PyTorch `preprocessors` the primary choice for training pipelines.
+- **Pure PyTorch Filtering** ✅: `BandPass` now uses a zero-phase FFT-based filter and `BaselineRemove` uses dual `avg_pool1d`, both implemented in `utils/utils_signal_t.py`. No more CPU-GPU data transfers.
+- **Unification of Managers** ⬜: Refactor `PreprocManager` and `AugmenterManager` to share a common base or registry, as their logic for managing sequences of transforms is very similar.
+- **Dimension Agnostic Transforms** ⬜ (augmenters pending): Preprocessors now handle arbitrary leading batch dimensions (`..., n_leads, siglen`). Augmenters still need to be updated.
+- **Numpy Version Maintenance** ✅: `_preprocessors` (NumPy) is kept only for offline data preparation or deployment environments without PyTorch, while making the PyTorch `preprocessors` the primary choice for training pipelines.
 
-### 5. Pandas 3.0 Migration and Dtype Consistency
+### 5. Pandas 3.0 Migration and Dtype Consistency ⬜ Not started
 **Objective**: Ensure compatibility with Pandas 3.0+, particularly concerning Arrow-backed strings and stricter type checking.
 - **Explicit Object Dtypes**: Explicitly set `dtype=object` for DataFrame columns intended to hold list-like objects (e.g., diagnoses, available signals) to prevent errors when Arrow-backed strings are used.
 - **Initialization Refactoring**: Replace patterns of initializing columns with `None` or `""` and then populating with lists via `.at[]` with more robust initializations like `[[] for _ in range(len(df))]` or using `.apply()`.
 - **Vectorized Operations**: Prefer `.apply()` or other vectorized pandas operations over `iterrows()` loops for better performance and type consistency.
 
 ---
 
-## Self-Supervised Learning (SSL) Roadmap
+## Complete Incomplete Models ⬜
 
+Several model files in `torch_ecg/models/` are currently stubs (`raise NotImplementedError` throughout). All stubs already inherit the correct mixins (`SizeMixin`, `CitationMixin`) and have backbone API signatures (`forward_features`, `compute_features_output_shape`) scaffolded. These should be completed before the SSL phase.
+
+| File | Classes | Notes |
+|---|---|---|
+| `models/cnn/darknet.py` | `DarkNet` | Backbone for YOLO-style detection |
+| `models/cnn/efficientnet.py` | `EfficientNet`, `EfficientNetV2` | Mobile-efficient compound scaling |
+| `models/cnn/ho_resnet.py` | `MidPointResNet`, `RK4ResNet`, `RK8ResNet` | Higher-Order ODE ResNets |
+| `models/grad_cam.py` | `GradCam` | Saliency analysis (logic partially drafted, not wired up) |
+| `models/ecg_fcn.py` | `ECG_FCN` | Fully-convolutional segmentation (fully commented out) |
+
+---
+
+## Self-Supervised Learning (SSL) Roadmap ⬜
+
+The `torch_ecg/ssl/` module is an empty shell. `ssl/README.md` contains a survey of target architectures (CLOCS, ST-MEM, MAE-ECG, SimCLR, TF-C, 3M-ECG, CMSC, ECG-BERT).
+
+**Prerequisites**: Item 3 (Lazy Modules) and the model stubs above should be completed first.
+
+- [ ] Define base classes: `BaseContrastiveLearner`, `BaseMaskedAutoencoder` (in `ssl/base.py`).
 - [ ] Implement a base contrastive learning framework (supporting paradigms like SimCLR, MoCo).
 - [ ] Implement Masked Autoencoder (MAE) logic optimized for 1D physiological signals.
 - [ ] Official implementation of classic models: CLOCS, ST-MEM, MAE-ECG.
 - [ ] Provide a unified Fine-tuning API for seamless integration with downstream tasks in `torch_ecg.models`.
+- [ ] Add `SSL` registry (analogous to `MODELS`, `BACKBONES`).
 
 ---
 

test/test_preproc_perf.py

@@ -0,0 +1,119 @@
+"""
+Performance and Numerical Consistency Benchmarking for Preprocessors.
+"""
+
+import time
+
+import numpy as np
+import pytest
+import torch
+
+from torch_ecg.preprocessors import BandPass, BaselineRemove
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def generate_test_signal(batch_size=32, n_leads=12, seq_len=5000, fs=500):
+    # Sinusoidal signal + baseline drift + noise
+    t = torch.linspace(0, seq_len / fs, seq_len)
+    # Pure signal (10 Hz)
+    sig = torch.sin(2 * np.pi * 10 * t)
+    # Baseline drift (0.1 Hz)
+    drift = 2.0 * torch.sin(2 * np.pi * 0.1 * t)
+    # Noise
+    noise = 0.5 * torch.randn(seq_len)
+
+    base_sig = sig + drift + noise
+    # Expand to batch and leads
+    batch_sig = base_sig.unsqueeze(0).unsqueeze(0).repeat(batch_size, n_leads, 1)
+    return batch_sig
+
+
+def test_bandpass_consistency():
+    fs = 500
+    lowcut, highcut = 0.5, 45
+    sig_torch = generate_test_signal(batch_size=4, seq_len=10000, fs=fs).to(DEVICE)
+    sig_np = sig_torch.cpu().numpy()
+
+    # 1. SciPy version (via legacy path)
+    bp = BandPass(fs=fs, lowcut=lowcut, highcut=highcut)
+    out_np = bp._forward_numpy(sig_np)
+
+    # 2. Torch version
+    out_torch = bp._forward_torch(sig_torch)
+
+    # Check consistency (FFT filter and IIR filter won't be identical, but should be close in passband)
+    # We use MSE as a loose metric because frequency domain zeroing != time domain recursion
+    diff = np.abs(out_np - out_torch.cpu().numpy())
+    mse = np.mean(diff**2)
+    print(f"\nBandPass Consistency (MSE): {mse:.6f}")
+    # FFT filtering is usually "cleaner" than IIR, so some difference is expected
+    assert mse < 0.1
+
+
+def test_baseline_remove_consistency():
+    fs = 500
+    window1, window2 = 0.2, 0.6
+    sig_torch = generate_test_signal(batch_size=4, seq_len=10000, fs=fs).to(DEVICE)
+    sig_np = sig_torch.cpu().numpy()
+
+    br = BaselineRemove(fs=fs, window1=window1, window2=window2)
+
+    # 1. SciPy version (Median Filter)
+    out_np = br._forward_numpy(sig_np)
+
+    # 2. Torch version (Sliding Average)
+    out_torch = br._forward_torch(sig_torch)
+
+    diff = np.abs(out_np - out_torch.cpu().numpy())
+    mse = np.mean(diff**2)
+    print(f"\nBaselineRemove Consistency (MSE): {mse:.6f}")
+    # Median filter and Sliding average are different algorithms, but should achieve same goal
+    assert mse < 0.5
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU performance test requires CUDA")
+def test_preproc_performance():
+    batch_size = 128
+    n_leads = 12
+    seq_len = 5000
+    fs = 500
+    sig_torch = generate_test_signal(batch_size, n_leads, seq_len, fs).to(DEVICE)
+
+    bp = BandPass(fs=fs, lowcut=0.5, highcut=45).to(DEVICE)
+    br = BaselineRemove(fs=fs).to(DEVICE)
+
+    # Warmup
+    _ = bp(sig_torch)
+    _ = br(sig_torch)
+    torch.cuda.synchronize()
+
+    # 1. Performance of Torch path
+    start = time.time()
+    for _ in range(10):
+        _ = bp._forward_torch(sig_torch)
+        _ = br._forward_torch(sig_torch)
+    torch.cuda.synchronize()
+    torch_time = (time.time() - start) / 10
+
+    # 2. Performance of NumPy path (includes data transfer simulation)
+    start = time.time()
+    for _ in range(10):
+        # Simulation of old way: move to CPU -> proc -> move back
+        tmp = sig_torch.cpu().numpy()
+        _ = bp._forward_numpy(tmp)
+        _ = br._forward_numpy(tmp)
+        _ = torch.as_tensor(tmp).to(DEVICE)
+    numpy_time = (time.time() - start) / 10
+
+    speedup = numpy_time / torch_time
+    print(f"\nPerformance Result (Batch size {batch_size}):")
+    print(f"NumPy Path (with transfers): {numpy_time:.4f}s")
+    print(f"Torch Path (on GPU): {torch_time:.4f}s")
+    print(f"Speedup: {speedup:.2f}x")
+
+    assert speedup > 2.0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

test/test_preprocessors_t.py

@@ -1,5 +1,6 @@
 """ """
 
+import numpy as np
 import pytest
 import torch
 
@@ -16,6 +17,7 @@
 )
 
 test_sig = torch.randn(2, 12, 8000)
+test_sig_np = test_sig.cpu().numpy()
 
 
 class DummyPreProcessor(torch.nn.Module):
@@ -40,8 +42,15 @@ def test_preproc_manager() -> None:
     ppm.add_(Normalize(method="min-max"))
     ppm.add_(Resample(fs=300, dst_fs=500), pos=0)
 
+    # Test with torch.Tensor
     sig = test_sig.clone()
     sig = ppm(sig)
+    assert isinstance(sig, torch.Tensor)
+
+    # Test with np.ndarray
+    sig_np = test_sig_np.copy()
+    sig_np = ppm(sig_np)
+    assert isinstance(sig_np, np.ndarray)
 
     config = CFG(
         random=False,
@@ -108,8 +117,14 @@ def test_preproc_manager() -> None:
 
 def test_bandpass() -> None:
     bp = BandPass(fs=500)
+    # Tensor
     sig = test_sig.clone()
     sig = bp(sig)
+    assert isinstance(sig, torch.Tensor)
+    # ndarray
+    sig_np = test_sig_np.copy()
+    sig_np = bp(sig_np)
+    assert isinstance(sig_np, np.ndarray)
 
     bp = BandPass(fs=500, lowcut=0, highcut=40)
     sig = test_sig.clone()
@@ -124,8 +139,14 @@ def test_bandpass() -> None:
 
 def test_baseline_remove() -> None:
     br = BaselineRemove(fs=500, inplace=False)
+    # Tensor
     sig = test_sig.clone()
     sig = br(sig)
+    assert isinstance(sig, torch.Tensor)
+    # ndarray
+    sig_np = test_sig_np.copy()
+    sig_np = br(sig_np)
+    assert isinstance(sig_np, np.ndarray)
 
     br = BaselineRemove(fs=500, window1=0.3, window2=0.7)
     sig = test_sig.clone()
@@ -139,8 +160,14 @@ def test_baseline_remove() -> None:
 
 def test_normalize() -> None:
     norm = Normalize(method="min-max", inplace=False)
+    # Tensor
     sig = test_sig.clone()
     sig = norm(sig)
+    assert isinstance(sig, torch.Tensor)
+    # ndarray
+    sig_np = test_sig_np.copy()
+    sig_np = norm(sig_np)
+    assert isinstance(sig_np, np.ndarray)
 
     norm = Normalize(method="z-score")
     sig = test_sig.clone()
@@ -159,3 +186,17 @@ def test_normalize() -> None:
     sig = norm(sig)
 
     del norm, sig
+
+
+def test_resample() -> None:
+    rsmp = Resample(fs=500, dst_fs=300)
+    # Tensor
+    sig = test_sig.clone()
+    sig = rsmp(sig)
+    assert isinstance(sig, torch.Tensor)
+    # ndarray
+    sig_np = test_sig_np.copy()
+    sig_np = rsmp(sig_np)
+    assert isinstance(sig_np, np.ndarray)
+
+    del rsmp, sig