Project-MONAI / Project-MONAI/MONAI

[Feature Request] Implement "Depth Reduced Strategy" for Model Compression across MONAI Networks

Open
#9,050 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
8.7k
Forks
1.6k
Avg merge
5d 1h
Merged PRs (30d)
20

Description

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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the existing model-compression and quantization support in MONAI, then compare the network entry points named here: SegResNet, SwinUNETR, UNet, and DiNTS. Before implementation, clarify whether the project wants a wrapper, built-in prune methods, or training-aware reduction, and narrow the work to a proof of concept for one architecture. Done should include an agreed API and evidence that the selected network can be depth-reduced and fine-tuned.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.