MSK-Net: Multi-Scale Spatial KANs Enhanced U-Shaped Network For Explainable 3D Brain Tumor Segmentation
Official PyTorch implementation of MSK-Net for explainable 3D brain tumor segmentation.
MSK-Net is a novel 3D segmentation framework that integrates Kolmogorov-Arnold Networks (KANs) into a U-shaped architecture for accurate and interpretable brain tumor segmentation from MRI scans.
- Spatial KAN Block (SKB): Replaces conventional linear filters with tensor products of univariate B-spline functions
- KAN-Boosted Attention Module (KBAM): Learnable spline-based channel-spatial attention mechanism
- Cross-Scale Gating Module (CSGM): Adaptive multi-scale feature fusion with KAN refinement
- Optimized 3D KAN Implementation: Vectorized basis caching and spatial parameter sharing
- Python >= 3.8
- PyTorch >= 2.1.0
- CUDA >= 12.x
# Clone repository
git clone https://github.com/msk-net/msk_net.git
cd msk_net
# Create virtual environment
conda create -n msk_net python=3.9
conda activate msk_net
# Install dependencies
pip install -r requirements.txtDownload BraTS datasets and organize data structure:
/path/to/dataset/
├── train/
│ ├── patient_001/
│ │ ├── *_t1.nii.gz
│ │ ├── *_t1ce.nii.gz
│ │ ├── *_t2.nii.gz
│ │ ├── *_flair.nii.gz
│ │ └── *_seg.nii.gz
│ └── ...
└── test/
└── ...
Copy the template configuration file and fill in your parameters:
cp configs/config_template.yaml configs/my_config.yaml
# Edit configs/my_config.yaml with your settingspython train.py --config configs/my_config.yaml
# Resume from checkpoint
python train.py --config configs/my_config.yaml --resume checkpoints/latest.pthpython inference.py \
--config configs/my_config.yaml \
--checkpoint checkpoints/best.pth \
--ttamsk_net/
├── configs/
│ ├── config_template.yaml # Configuration template
│ └── ...
├── models/
│ ├── blocks/
│ │ ├── skb.py # Spatial KAN Block
│ │ ├── kbam.py # KAN-Boosted Attention Module
│ │ ├── csgm.py # Cross-Scale Gating Module
│ │ ├── aspp.py # Atrous Spatial Pyramid Pooling
│ │ └── common.py # Common layers
│ └── msk_net.py # Main architecture
├── datasets/
│ └── brats_dataset.py # Data loading and preprocessing
├── losses/
│ ├── losses.py # Loss functions
│ └── metrics.py # Evaluation metrics
├── explainability/
│ ├── score_cam.py # Score-CAM implementation
│ ├── grad_cam.py # Grad-CAM implementation
│ └── full_grad.py # FullGrad implementation
├── train.py # Training script
├── inference.py # Inference and testing
└── README.md
MSK-Net supports multiple KAN basis function implementations:
- B-Spline (default)
- Fourier
- Chebyshev
- Hermite
- Gegenbauer
- Jacobi
- Bessel
- Lucas
- Fibonacci
- Gaussian RBF
- Wavelet
This project is licensed under the MIT License.
- KAN implementations from xKAN
- BraTS challenge organizers