wavedl 1.6.2__tar.gz → 1.7.0__tar.gz

{wavedl-1.6.2/src/wavedl.egg-info → wavedl-1.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: wavedl
-Version: 1.6.2
+Version: 1.7.0
 Summary: A Scalable Deep Learning Framework for Wave-Based Inverse Problems
 Author: Ductho Le
 License: MIT
@@ -214,11 +214,11 @@ This installs everything you need: training, inference, HPO, ONNX export.
 ```bash
 git clone https://github.com/ductho-le/WaveDL.git
 cd WaveDL
-pip install -e .
+pip install -e ".[dev]"
 ```
 
 > [!NOTE]
-> Python 3.11+ required. For development setup, see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
+> Python 3.11+ required. For contributor setup (pre-commit hooks), see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 
 ### Quick Start
 
@@ -273,8 +273,6 @@ accelerate launch --num_machines 2 --main_process_ip <ip> -m wavedl.train --mode
 
 ### Testing & Inference
 
-After training, use `wavedl-test` to evaluate your model on test data:
-
 ```bash
 # Basic inference
 wavedl-test --checkpoint <checkpoint_folder> --data_path <test_data>
@@ -909,18 +907,24 @@ Automatically find the best training configuration using [Optuna](https://optuna
 **Run HPO:**
 
 ```bash
-# Basic HPO (auto-detects GPUs for parallel trials)
-wavedl-hpo --data_path train.npz --models cnn --n_trials 100
+# Basic HPO (50 trials, auto-detects GPUs)
+wavedl-hpo --data_path train.npz --n_trials 50
+
+# Quick search (minimal search space, fastest)
+wavedl-hpo --data_path train.npz --n_trials 30 --quick
 
-# Search multiple models
-wavedl-hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 200
+# Medium search (balanced between quick and full)
+wavedl-hpo --data_path train.npz --n_trials 50 --medium
 
-# Quick mode (fewer parameters, faster)
-wavedl-hpo --data_path train.npz --models cnn --n_trials 50 --quick
+# Full search with specific models
+wavedl-hpo --data_path train.npz --n_trials 100 --models cnn resnet18 efficientnet_b0
+
+# In-process mode (enables pruning, faster, single-GPU)
+wavedl-hpo --data_path train.npz --n_trials 50 --inprocess
 ```
 
 > [!TIP]
-> **Auto GPU Detection**: HPO automatically detects available GPUs and runs one trial per GPU in parallel. On a 4-GPU system, 4 trials run simultaneously. Use `--n_jobs 1` to force serial execution.
+> **GPU Detection**: HPO auto-detects GPUs and runs one trial per GPU in parallel. Use `--inprocess` for single-GPU with pruning support (early stopping of bad trials).
 
 **Train with best parameters**
 
@@ -942,10 +946,23 @@ wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...
 | Learning rate | 1e-5 → 1e-2 | (always searched) |
 | Batch size | 16, 32, 64, 128 | (always searched) |
 
-**Quick Mode** (`--quick`):
-- Uses minimal defaults: cnn + adamw + plateau + mse
-- Faster for testing your setup before running full search
-- You can still override any option with the flags above
+**Search Presets:**
+
+| Mode | Models | Optimizers | Schedulers | Use Case |
+|------|--------|------------|------------|----------|
+| Full (default) | cnn, resnet18, resnet34 | all 6 | all 8 | Production search |
+| `--medium` | cnn, resnet18 | adamw, adam, sgd | plateau, cosine, onecycle | Balanced exploration |
+| `--quick` | cnn | adamw | plateau | Fast validation |
+
+**Execution Modes:**
+
+| Mode | Flag | Pruning | GPU Memory | Best For |
+|------|------|---------|------------|----------|
+| Subprocess (default) | — | ❌ No | Isolated | Multi-GPU parallel trials |
+| In-process | `--inprocess` | ✅ Yes | Shared | Single-GPU with early stopping |
+
+> [!TIP]
+> Use `--inprocess` when running single-GPU trials. It enables MedianPruner to stop unpromising trials early, reducing total search time.
 
 ---
 
@@ -956,7 +973,9 @@ wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...
 | `--data_path` | (required) | Training data file |
 | `--models` | 3 defaults | Models to search (specify any number) |
 | `--n_trials` | `50` | Number of trials to run |
-| `--quick` | `False` | Use minimal defaults (faster) |
+| `--quick` | `False` | Quick mode: minimal search space |
+| `--medium` | `False` | Medium mode: balanced search space |
+| `--inprocess` | `False` | Run trials in-process (enables pruning) |
 | `--optimizers` | all 6 | Optimizers to search |
 | `--schedulers` | all 8 | Schedulers to search |
 | `--losses` | all 6 | Losses to search |
@@ -965,7 +984,7 @@ wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...
 | `--output` | `hpo_results.json` | Output file |
 
 
-> See [Available Models](#available-models) for all 38 architectures you can search.
+> See [Available Models](#available-models) for all 69 architectures you can search.
 
 </details>
 

{wavedl-1.6.2 → wavedl-1.7.0}/README.md

@@ -166,11 +166,11 @@ This installs everything you need: training, inference, HPO, ONNX export.
 ```bash
 git clone https://github.com/ductho-le/WaveDL.git
 cd WaveDL
-pip install -e .
+pip install -e ".[dev]"
 ```
 
 > [!NOTE]
-> Python 3.11+ required. For development setup, see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
+> Python 3.11+ required. For contributor setup (pre-commit hooks), see [CONTRIBUTING.md](.github/CONTRIBUTING.md).
 
 ### Quick Start
 
@@ -225,8 +225,6 @@ accelerate launch --num_machines 2 --main_process_ip <ip> -m wavedl.train --mode
 
 ### Testing & Inference
 
-After training, use `wavedl-test` to evaluate your model on test data:
-
 ```bash
 # Basic inference
 wavedl-test --checkpoint <checkpoint_folder> --data_path <test_data>
@@ -861,18 +859,24 @@ Automatically find the best training configuration using [Optuna](https://optuna
 **Run HPO:**
 
 ```bash
-# Basic HPO (auto-detects GPUs for parallel trials)
-wavedl-hpo --data_path train.npz --models cnn --n_trials 100
+# Basic HPO (50 trials, auto-detects GPUs)
+wavedl-hpo --data_path train.npz --n_trials 50
+
+# Quick search (minimal search space, fastest)
+wavedl-hpo --data_path train.npz --n_trials 30 --quick
 
-# Search multiple models
-wavedl-hpo --data_path train.npz --models cnn resnet18 efficientnet_b0 --n_trials 200
+# Medium search (balanced between quick and full)
+wavedl-hpo --data_path train.npz --n_trials 50 --medium
 
-# Quick mode (fewer parameters, faster)
-wavedl-hpo --data_path train.npz --models cnn --n_trials 50 --quick
+# Full search with specific models
+wavedl-hpo --data_path train.npz --n_trials 100 --models cnn resnet18 efficientnet_b0
+
+# In-process mode (enables pruning, faster, single-GPU)
+wavedl-hpo --data_path train.npz --n_trials 50 --inprocess
 ```
 
 > [!TIP]
-> **Auto GPU Detection**: HPO automatically detects available GPUs and runs one trial per GPU in parallel. On a 4-GPU system, 4 trials run simultaneously. Use `--n_jobs 1` to force serial execution.
+> **GPU Detection**: HPO auto-detects GPUs and runs one trial per GPU in parallel. Use `--inprocess` for single-GPU with pruning support (early stopping of bad trials).
 
 **Train with best parameters**
 
@@ -894,10 +898,23 @@ wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...
 | Learning rate | 1e-5 → 1e-2 | (always searched) |
 | Batch size | 16, 32, 64, 128 | (always searched) |
 
-**Quick Mode** (`--quick`):
-- Uses minimal defaults: cnn + adamw + plateau + mse
-- Faster for testing your setup before running full search
-- You can still override any option with the flags above
+**Search Presets:**
+
+| Mode | Models | Optimizers | Schedulers | Use Case |
+|------|--------|------------|------------|----------|
+| Full (default) | cnn, resnet18, resnet34 | all 6 | all 8 | Production search |
+| `--medium` | cnn, resnet18 | adamw, adam, sgd | plateau, cosine, onecycle | Balanced exploration |
+| `--quick` | cnn | adamw | plateau | Fast validation |
+
+**Execution Modes:**
+
+| Mode | Flag | Pruning | GPU Memory | Best For |
+|------|------|---------|------------|----------|
+| Subprocess (default) | — | ❌ No | Isolated | Multi-GPU parallel trials |
+| In-process | `--inprocess` | ✅ Yes | Shared | Single-GPU with early stopping |
+
+> [!TIP]
+> Use `--inprocess` when running single-GPU trials. It enables MedianPruner to stop unpromising trials early, reducing total search time.
 
 ---
 
@@ -908,7 +925,9 @@ wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...
 | `--data_path` | (required) | Training data file |
 | `--models` | 3 defaults | Models to search (specify any number) |
 | `--n_trials` | `50` | Number of trials to run |
-| `--quick` | `False` | Use minimal defaults (faster) |
+| `--quick` | `False` | Quick mode: minimal search space |
+| `--medium` | `False` | Medium mode: balanced search space |
+| `--inprocess` | `False` | Run trials in-process (enables pruning) |
 | `--optimizers` | all 6 | Optimizers to search |
 | `--schedulers` | all 8 | Schedulers to search |
 | `--losses` | all 6 | Losses to search |
@@ -917,7 +936,7 @@ wavedl-train --data_path train.npz --model cnn --lr 3.2e-4 --batch_size 128 ...
 | `--output` | `hpo_results.json` | Output file |
 
 
-> See [Available Models](#available-models) for all 38 architectures you can search.
+> See [Available Models](#available-models) for all 69 architectures you can search.
 
 </details>
 

{wavedl-1.6.2 → wavedl-1.7.0}/src/wavedl/__init__.py

@@ -18,7 +18,7 @@ For inference:
     # or: python -m wavedl.test --checkpoint best_checkpoint --data_path test.npz
 """
 
-__version__ = "1.6.2"
+__version__ = "1.7.0"
 __author__ = "Ductho Le"
 __email__ = "ductho.le@outlook.com"
 

{wavedl-1.6.2 → wavedl-1.7.0}/src/wavedl/hpo.py

@@ -10,12 +10,28 @@ Usage:
     # Quick search (fewer parameters)
     wavedl-hpo --data_path train.npz --n_trials 30 --quick
 
+    # Medium search (balanced)
+    wavedl-hpo --data_path train.npz --n_trials 50 --medium
+
     # Full search with specific models
     wavedl-hpo --data_path train.npz --n_trials 100 --models cnn resnet18 efficientnet_b0
 
     # Parallel trials on multiple GPUs
     wavedl-hpo --data_path train.npz --n_trials 100 --n_jobs 4
 
+    # In-process mode (enables pruning, faster, single-GPU)
+    wavedl-hpo --data_path train.npz --n_trials 50 --inprocess
+
+Execution Modes:
+    --inprocess: Runs trials in the same Python process. Enables pruning
+                 (MedianPruner) for early stopping of unpromising trials.
+                 Faster due to no subprocess overhead, but trials share
+                 GPU memory (no isolation between trials).
+
+    Default (subprocess): Launches each trial as a separate process.
+                          Provides GPU memory isolation but prevents pruning
+                          (subprocess can't report intermediate results).
+
 Author: Ductho Le (ductho.le@outlook.com)
 """
 
@@ -41,10 +57,12 @@ except ImportError:
 
 DEFAULT_MODELS = ["cnn", "resnet18", "resnet34"]
 QUICK_MODELS = ["cnn"]
+MEDIUM_MODELS = ["cnn", "resnet18"]
 
 # All 6 optimizers
 DEFAULT_OPTIMIZERS = ["adamw", "adam", "sgd", "nadam", "radam", "rmsprop"]
 QUICK_OPTIMIZERS = ["adamw"]
+MEDIUM_OPTIMIZERS = ["adamw", "adam", "sgd"]
 
 # All 8 schedulers
 DEFAULT_SCHEDULERS = [
@@ -58,10 +76,12 @@ DEFAULT_SCHEDULERS = [
     "linear_warmup",
 ]
 QUICK_SCHEDULERS = ["plateau"]
+MEDIUM_SCHEDULERS = ["plateau", "cosine", "onecycle"]
 
 # All 6 losses
 DEFAULT_LOSSES = ["mse", "mae", "huber", "smooth_l1", "log_cosh", "weighted_mse"]
 QUICK_LOSSES = ["mse"]
+MEDIUM_LOSSES = ["mse", "mae", "huber"]
 
 
 # =============================================================================
@@ -70,16 +90,28 @@ QUICK_LOSSES = ["mse"]
 
 
 def create_objective(args):
-    """Create Optuna objective function with configurable search space."""
+    """Create Optuna objective function with configurable search space.
+
+    Supports two execution modes:
+    - Subprocess (default): Launches wavedl.train via subprocess. Provides GPU
+      memory isolation but prevents pruning (MedianPruner has no effect).
+    - In-process (--inprocess): Calls train_single_trial() directly. Enables
+      pruning and reduces overhead, but trials share GPU memory.
+    """
 
     def objective(trial):
-        # Select search space based on mode
+        # Select search space based on mode (quick < medium < full)
         # CLI arguments always take precedence over defaults
         if args.quick:
             models = args.models or QUICK_MODELS
             optimizers = args.optimizers or QUICK_OPTIMIZERS
             schedulers = args.schedulers or QUICK_SCHEDULERS
             losses = args.losses or QUICK_LOSSES
+        elif args.medium:
+            models = args.models or MEDIUM_MODELS
+            optimizers = args.optimizers or MEDIUM_OPTIMIZERS
+            schedulers = args.schedulers or MEDIUM_SCHEDULERS
+            losses = args.losses or MEDIUM_LOSSES
         else:
             models = args.models or DEFAULT_MODELS
             optimizers = args.optimizers or DEFAULT_OPTIMIZERS
@@ -101,13 +133,59 @@ def create_objective(args):
         if loss == "huber":
             huber_delta = trial.suggest_float("huber_delta", 0.1, 2.0)
         else:
-            huber_delta = None
+            huber_delta = 1.0  # default
 
         if optimizer == "sgd":
             momentum = trial.suggest_float("momentum", 0.8, 0.99)
         else:
-            momentum = None
+            momentum = 0.9  # default
+
+        # ==================================================================
+        # IN-PROCESS MODE: Direct function call with pruning support
+        # ==================================================================
+        if args.inprocess:
+            from wavedl.train import train_single_trial
+
+            try:
+                result = train_single_trial(
+                    data_path=args.data_path,
+                    model_name=model,
+                    lr=lr,
+                    batch_size=batch_size,
+                    epochs=args.max_epochs,
+                    patience=patience,
+                    optimizer_name=optimizer,
+                    scheduler_name=scheduler,
+                    loss_name=loss,
+                    weight_decay=weight_decay,
+                    seed=args.seed,
+                    huber_delta=huber_delta,
+                    momentum=momentum,
+                    trial=trial,  # Enable pruning via trial.report/should_prune
+                    verbose=False,
+                )
+
+                if result["pruned"]:
+                    print(
+                        f"Trial {trial.number}: Pruned at epoch {result['epochs_trained']}"
+                    )
+                    raise optuna.TrialPruned()
+
+                val_loss = result["best_val_loss"]
+                print(
+                    f"Trial {trial.number}: val_loss={val_loss:.6f} ({result['epochs_trained']} epochs)"
+                )
+                return val_loss
 
+            except optuna.TrialPruned:
+                raise  # Re-raise for Optuna to handle
+            except Exception as e:
+                print(f"Trial {trial.number}: Error - {e}")
+                return float("inf")
+
+        # ==================================================================
+        # SUBPROCESS MODE (default): GPU memory isolation, no pruning
+        # ==================================================================
         # Build command
         cmd = [
             sys.executable,
@@ -138,9 +216,9 @@ def create_objective(args):
         ]
 
         # Add conditional args
-        if huber_delta:
+        if loss == "huber":
             cmd.extend(["--huber_delta", str(huber_delta)])
-        if momentum:
+        if optimizer == "sgd":
             cmd.extend(["--momentum", str(momentum)])
 
         # Use temporary directory for trial output
@@ -285,7 +363,17 @@ Examples:
     parser.add_argument(
         "--quick",
         action="store_true",
-        help="Quick mode: search fewer parameters",
+        help="Quick mode: search fewer parameters (fastest, least thorough)",
+    )
+    parser.add_argument(
+        "--medium",
+        action="store_true",
+        help="Medium mode: balanced parameter search (between --quick and full)",
+    )
+    parser.add_argument(
+        "--inprocess",
+        action="store_true",
+        help="Run trials in-process (enables pruning, faster, but no GPU memory isolation)",
     )
     parser.add_argument(
         "--timeout",
@@ -384,14 +472,32 @@ Examples:
     print("=" * 60)
     print(f"Data: {args.data_path}")
     print(f"Trials: {args.n_trials}")
-    print(f"Mode: {'Quick' if args.quick else 'Full'}")
+    # Determine mode name for display
+    if args.quick:
+        mode_name = "Quick"
+    elif args.medium:
+        mode_name = "Medium"
+    else:
+        mode_name = "Full"
+
+    print(
+        f"Mode: {mode_name}"
+        + (" (in-process, pruning enabled)" if args.inprocess else " (subprocess)")
+    )
     print(f"Parallel jobs: {args.n_jobs}")
     print("=" * 60)
 
+    # Use MedianPruner only for in-process mode (subprocess trials can't report)
+    if args.inprocess:
+        pruner = optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=10)
+    else:
+        # NopPruner for subprocess mode - pruning has no effect there
+        pruner = optuna.pruners.NopPruner()
+
     study = optuna.create_study(
         study_name=args.study_name,
         direction="minimize",
-        pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=10),
+        pruner=pruner,
     )
 
     # Run optimization

{wavedl-1.6.2 → wavedl-1.7.0}/src/wavedl/models/__init__.py

@@ -77,6 +77,15 @@ from .unet import UNetRegression
 from .vit import ViTBase_, ViTSmall, ViTTiny
 
 
+# Optional RATENet (unpublished, may be gitignored)
+try:
+    from .ratenet import RATENet, RATENetLite, RATENetTiny, RATENetV2
+
+    _HAS_RATENET = True
+except ImportError:
+    _HAS_RATENET = False
+
+
 # Optional timm-based models (imported conditionally)
 try:
     from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
@@ -111,6 +120,7 @@ __all__ = [
     "MC3_18",
     "MODEL_REGISTRY",
     "TCN",
+    # Classes (uppercase first, alphabetically)
     "BaseModel",
     "ConvNeXtBase_",
     "ConvNeXtSmall",
@@ -152,6 +162,7 @@ __all__ = [
     "VimBase",
     "VimSmall",
     "VimTiny",
+    # Functions (lowercase, alphabetically)
     "build_model",
     "get_model",
     "list_models",
@@ -186,3 +197,14 @@ if _HAS_TIMM_MODELS:
             "UniRepLKNetTiny",
         ]
     )
+
+# Add RATENet models to __all__ if available (unpublished)
+if _HAS_RATENET:
+    __all__.extend(
+        [
+            "RATENet",
+            "RATENetLite",
+            "RATENetTiny",
+            "RATENetV2",
+        ]
+    )

{wavedl-1.6.2 → wavedl-1.7.0}/src/wavedl/models/_pretrained_utils.py

@@ -166,6 +166,78 @@ class LayerNormNd(nn.Module):
         return x
 
 
+# =============================================================================
+# STOCHASTIC DEPTH (DropPath)
+# =============================================================================
+
+
+class DropPath(nn.Module):
+    """
+    Stochastic Depth (drop path) regularization for residual networks.
+
+    Randomly drops entire residual branches during training. Used in modern
+    architectures like ConvNeXt, Swin Transformer, UniRepLKNet.
+
+    Args:
+        drop_prob: Probability of dropping the path (default: 0.0)
+
+    Reference:
+        Huang, G., et al. (2016). Deep Networks with Stochastic Depth.
+        https://arxiv.org/abs/1603.09382
+    """
+
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+
+        keep_prob = 1 - self.drop_prob
+        # Shape: (batch_size, 1, 1, ...) for broadcasting
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()  # Binarize
+        return x.div(keep_prob) * random_tensor
+
+
+# =============================================================================
+# BACKBONE FREEZING UTILITIES
+# =============================================================================
+
+
+def freeze_backbone(
+    model: nn.Module,
+    exclude_patterns: list[str] | None = None,
+) -> int:
+    """
+    Freeze backbone parameters, keeping specified layers trainable.
+
+    Args:
+        model: The model whose parameters to freeze
+        exclude_patterns: List of patterns to exclude from freezing.
+            Parameters with names containing any of these patterns stay trainable.
+            Default: ["classifier", "head", "fc"]
+
+    Returns:
+        Number of parameters frozen
+
+    Example:
+        >>> freeze_backbone(model.backbone, exclude_patterns=["fc", "classifier"])
+    """
+    if exclude_patterns is None:
+        exclude_patterns = ["classifier", "head", "fc"]
+
+    frozen_count = 0
+    for name, param in model.named_parameters():
+        if not any(pattern in name for pattern in exclude_patterns):
+            param.requires_grad = False
+            frozen_count += param.numel()
+
+    return frozen_count
+
+
 # =============================================================================
 # REGRESSION HEAD BUILDERS
 # =============================================================================

{wavedl-1.6.2 → wavedl-1.7.0}/src/wavedl/models/_template.py

@@ -31,22 +31,23 @@ from wavedl.models.base import BaseModel
 # @register_model("my_model")
 class TemplateModel(BaseModel):
     """
-    Template Model Architecture.
+    Template Model Architecture (2D only).
 
     Replace this docstring with your model description.
     The first line will appear in --list_models output.
 
+    NOTE: This template is hardcoded for 2D inputs using Conv2d/MaxPool2d.
+    For 1D/3D support, use dimension-agnostic layer factories from
+    _pretrained_utils.py (get_conv_layer, get_pool_layer, get_norm_layer).
+
     Args:
-        in_shape: Input spatial dimensions (auto-detected from data)
-                  - 1D: (L,) for signals
-                  - 2D: (H, W) for images
-                  - 3D: (D, H, W) for volumes
+        in_shape: Input spatial dimensions as (H, W) for 2D images
         out_size: Number of regression targets (auto-detected from data)
         hidden_dim: Size of hidden layers (default: 256)
         dropout: Dropout rate (default: 0.1)
 
     Input Shape:
-        (B, 1, *in_shape) - e.g., (B, 1, 64, 64) for 2D
+        (B, 1, H, W) - 2D grayscale images
 
     Output Shape:
         (B, out_size) - Regression predictions

{wavedl-1.6.2 → wavedl-1.7.0}/src/wavedl/models/cnn.py

@@ -159,6 +159,26 @@ class CNN(BaseModel):
             nn.Linear(64, out_size),
         )
 
+        # Initialize weights
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize weights with Kaiming for conv, Xavier for linear."""
+        for m in self.modules():
+            if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+                nn.init.kaiming_normal_(
+                    m.weight, mode="fan_out", nonlinearity="leaky_relu"
+                )
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, (nn.GroupNorm, nn.LayerNorm)):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+
     def _make_conv_block(
         self, in_channels: int, out_channels: int, dropout: float = 0.0
     ) -> nn.Sequential: