Currently, MONAI provides state-of-the-art medical image segmentation and classification networks (e.g., SegResNet, SwinUNETR, UNet, DiNTS). However, for deployment on edge devices, clinical workstations with limited GPU memory, or for faster inference, model compression is highly necessary.
I propose introducing a Depth Reduced Strategy—a structured pruning technique across MONAI's core network architectures. This strategy would dynamically reduce the depth (number of layers or blocks) of a model during training or fine-tuning, effectively dropping redundant layers without requiring an entirely new architectural definition.
Motivation:
Resource Constraints: Medical imaging models inherently deal with 3D data, which consumes massive amounts of VRAM. A depth-reduced strategy would lower the memory footprint.
Inference Speed: Reducing network depth directly reduces latency, which is critical for real-time surgical pipelines or large-scale clinical batch processing.
Standardization: While users can manually build smaller networks by tweaking parameters like num_blocks or depth, a standardized depth reduction/pruning wrapper or utility across monai.networks would provide a unified API for model compression.
Proposed Implementation / Design:
The "Depth Reduced Strategy" could be implemented in a few potential ways:
Model Wrapper: A PyTorch wrapper class (e.g., monai.networks.utils.DepthPruner) that takes an existing nn.Module and systematically removes specific layers/blocks based on an importance criterion (e.g., L1-norm of weights or Taylor expansion) or a predefined reduction ratio.
Built-in prune Methods: Extending the base classes of popular networks (like SegResNet or UNet) with a .prune(depth_reduction_ratio: float) method, which restructures the nn.Sequential blocks dynamically.
Training-Aware Depth Reduction: Incorporating techniques like Stochastic Depth (DropPath) combined with magnitude-based pruning so the network learns to bypass layers that are eventually removed.
Alternatives Considered:
Post-training quantization (PTQ) or Quantization-Aware Training (QAT): MONAI already has some support for this, but quantization addresses precision, not architectural depth/FLOPs.
Existing parameter tweaks (e.g., passing smaller depths to SwinUNETR): This requires retraining from scratch. A depth-reduced strategy could ideally be applied to fine-tune pre-trained models from the MONAI Model Zoo.
Call to Action:
I would love to hear the maintainers' thoughts on whether this fits within MONAI's roadmap for deployment and model optimization. If this aligns with the project goals, I am happy to help contribute a proof-of-concept for a specific architecture (e.g., SegResNet) as a starting point.
Currently, MONAI provides state-of-the-art medical image segmentation and classification networks (e.g., SegResNet, SwinUNETR, UNet, DiNTS). However, for deployment on edge devices, clinical workstations with limited GPU memory, or for faster inference, model compression is highly necessary.
I propose introducing a Depth Reduced Strategy—a structured pruning technique across MONAI's core network architectures. This strategy would dynamically reduce the depth (number of layers or blocks) of a model during training or fine-tuning, effectively dropping redundant layers without requiring an entirely new architectural definition.
Motivation:
Resource Constraints: Medical imaging models inherently deal with 3D data, which consumes massive amounts of VRAM. A depth-reduced strategy would lower the memory footprint.
Inference Speed: Reducing network depth directly reduces latency, which is critical for real-time surgical pipelines or large-scale clinical batch processing.
Standardization: While users can manually build smaller networks by tweaking parameters like num_blocks or depth, a standardized depth reduction/pruning wrapper or utility across monai.networks would provide a unified API for model compression.
Proposed Implementation / Design:
The "Depth Reduced Strategy" could be implemented in a few potential ways:
Model Wrapper: A PyTorch wrapper class (e.g., monai.networks.utils.DepthPruner) that takes an existing nn.Module and systematically removes specific layers/blocks based on an importance criterion (e.g., L1-norm of weights or Taylor expansion) or a predefined reduction ratio.
Built-in prune Methods: Extending the base classes of popular networks (like SegResNet or UNet) with a .prune(depth_reduction_ratio: float) method, which restructures the nn.Sequential blocks dynamically.
Training-Aware Depth Reduction: Incorporating techniques like Stochastic Depth (DropPath) combined with magnitude-based pruning so the network learns to bypass layers that are eventually removed.
Alternatives Considered:
Post-training quantization (PTQ) or Quantization-Aware Training (QAT): MONAI already has some support for this, but quantization addresses precision, not architectural depth/FLOPs.
Existing parameter tweaks (e.g., passing smaller depths to SwinUNETR): This requires retraining from scratch. A depth-reduced strategy could ideally be applied to fine-tune pre-trained models from the MONAI Model Zoo.
Call to Action:
I would love to hear the maintainers' thoughts on whether this fits within MONAI's roadmap for deployment and model optimization. If this aligns with the project goals, I am happy to help contribute a proof-of-concept for a specific architecture (e.g., SegResNet) as a starting point.