Score-based Diffusion MRI Reconstruction (HPC-Modified)

Modified implementation of score-based diffusion models for accelerated MRI reconstruction, adapted for HPC training and optimized for multi-coil FastMRI data. Based on the original score-MRI repository.

Original Paper: Score-based diffusion models for accelerated MRI
Hyungjin Chung and Jong Chul Ye
Medical Image Analysis 2022

Note on Directory Structure

Training configurations and scripts used in this implementation are located in train_score/. This includes:

• Training configuration: train_score/config.py
• Training scripts: train_score/train.sh, train_score/eval.sh, train_score/recon.sh
• SLURM batch job: train_score/score.sbatch
• Segmentation fault fix documentation: train_score/SEGFAULT_FIX_LOG.md

Key Modifications & Improvements

This repository extends the original score-MRI codebase with the following improvements:

🚀 HPC Training Support

Successfully adapted the original codebase to run on high-performance computing environments:

• TensorFlow-PyTorch Conflict Resolution: The original repository used TensorFlow for file I/O while the model training was implemented in PyTorch, causing CUDA initialization conflicts and segmentation faults. We introduced minimal mock TensorFlow interfaces to bypass TF GPU initialization, allowing PyTorch to retain exclusive GPU access while preserving all file-loading functionality. See SEGFAULT_FIX_LOG.md for detailed documentation.

• Environment Configuration: Added SLURM batch job scripts (score.sbatch) with proper CUDA memory management settings:

  export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
  export CUDA_LAUNCH_BLOCKING=1

• Training Stability Improvements: Enhanced training configuration with:

  • Likelihood weighting for stable continuous-time training
  • Adaptive learning rate (1e-4) optimized for batch size 1
  • Configurable snapshot sampling frequency to reduce computational overhead
  • Data centering to [-1, 1] range for better score model training

• Singularity Container Support: Configured for containerized execution on HPC systems with GPU access.

📊 FastMRI Multi-Coil Dataset Compatibility

Enhanced data handling for clinical multi-coil datasets:

• Coil Normalization: Fixed coil dimensions to 15 coils across all samples for batch consistency
• Image Resizing: Automatic resizing/padding to 320×320 resolution
• HDF5 File Support: Direct loading from FastMRI .h5 files with slice selection
• Robust Data Processing:
  • Explicit k-space to image domain conversion tracking
  • Percentile-based normalization (99th percentile) for stability across samples
  • Proper complex data handling with magnitude extraction
  • Automatic RSS (Root Sum of Squares) computation for multi-coil data

⚡ SENSE-Based Inference (15× Speedup)

Major Performance Improvement: Implemented SENSE (SENSitivity Encoding) parallel imaging reconstruction following Algorithm 3 from the paper.

Speed Comparison:

• Original (per-coil): ~30,000 model evaluations (15 coils × 2,000 steps)
• SENSE (combined): ~2,000 model evaluations (1 combined image × 2,000 steps)
• Result: 15× faster inference 🎯

Technical Implementation:

• Automatic coil sensitivity map estimation using ACS (Auto-Calibration Signal) extraction
• Low-resolution reconstruction with Gaussian smoothing
• Single combined image reconstruction instead of per-coil processing
• Full compatibility with the original NCSN++ model architecture

See inference_multi-coil_SENSE_h5.py for the optimized SENSE implementation.

📝 Additional Utilities

• Pure Generation: pure_gen_sense.py for generating random knee MRI images from noise without data consistency (pure generative sampling)
• Batch Reconstruction: batch_reconstruct_val.py for processing multiple validation samples
• Sensitivity Map Comparison: compare_sensitivity_methods.py for evaluating different estimation methods
• Metrics Tracking: Built-in MSE/PSNR calculation and CSV logging during reconstruction progress
• Training Scripts: Configured shell scripts for training (train.sh), evaluation (eval.sh), and reconstruction (recon.sh)
• Checkpoint Sweeping: sweep_sense_checkpoints_equi_acc4.py for evaluating multiple checkpoints
• Configuration Management: Custom config files adapted for HPC paths and settings

Installation

source install.sh

The installation script handles downloading model weights and installing dependencies.

Alternatively, download model weights here and place as weights/checkpoint_95.pth.

Training on HPC

Using SLURM

sbatch score.sbatch

Or run directly:

./train.sh

The training script is configured for:

• FastMRI knee multi-coil dataset
• 320×320 image resolution
• 15 coil channels
• VESDE (Variance Exploding SDE)
• NCSN++ architecture

Configuration

Modify config.py or your config file to adjust:

• Dataset paths
• Model hyperparameters
• Training settings
• Device selection

Inference

Multi-Coil SENSE Reconstruction (Recommended - 15× Faster)

python inference_multi-coil_SENSE_h5.py \
    --data /path/to/fastmri_file.h5 \
    --checkpoint /path/to/checkpoint.pth \
    --N 500 \
    --m 1 \
    --acc_factor 4 \
    --center_fraction 0.08 \
    --save_dir ./results

Performance: Reconstructs a 320×320 image in ~60 seconds (vs. ~900 seconds with per-coil method)

Pure Generation (No Data Consistency)

Generate random knee MRI images from pure noise:

python pure_gen_sense.py \
    --checkpoint /path/to/checkpoint.pth \
    --N 500 \
    --m 1 \
    --save_dir ./results

This demonstrates the generative capabilities of the trained score model without any reconstruction constraints.

Multi-Coil SSOS Reconstruction

python inference_multi-coil_SSOS_h5.py \
    --data /path/to/fastmri_file.h5 \
    --checkpoint /path/to/checkpoint.pth \
    --N 500 \
    --m 1 \
    --acc_factor 4 \
    --save_dir ./results

Hybrid Reconstruction

python inference_multi-coil_hybrid.py \
    --data /path/to/data.npy \
    --checkpoint /path/to/checkpoint.pth \
    --N 2000

Algorithm Correspondence

• inference_single-coil.py: Algorithm 2 from paper
• inference_multi-coil_SSOS.py: Algorithm 4 from paper
• inference_multi-coil_SSOS_h5.py: Algorithm 4 with HDF5 support and metrics logging
• inference_multi-coil_SENSE_h5.py: Algorithm 3 from paper (SENSE - 15× speedup) ⚡
• inference_multi-coil_hybrid.py: Algorithm 5 from paper
• pure_gen_sense.py: Pure generation (no data consistency)

Project Structure

├── configs/                    # Model and training configurations
│   ├── default_lsun_configs.py
│   └── ve/
│       └── fastmri_knee_320_ncsnpp_continuous_multi.py
├── models/                     # NCSN++ model architecture
├── op/                         # CUDA kernels
├── train_score/                # HPC training scripts and configs
│   ├── config.py               # Training configuration with stability improvements
│   ├── train.sh                # Training script
│   ├── eval.sh                 # Evaluation script
│   ├── recon.sh                # Reconstruction script
│   ├── score.sbatch            # SLURM batch job script
│   ├── SEGFAULT_FIX_LOG.md     # Documentation of TF/PyTorch conflict resolution
│   └── sweep_sense_checkpoints_equi_acc4.py  # Multi-checkpoint evaluation
├── utils.py                    # Data processing utilities (with sensitivity map estimation)
├── fastmri_utils.py            # MRI-specific transformations
├── datasets.py                 # Enhanced FastMRI data loading with robust normalization
├── sampling.py                 # Reconstruction algorithms with metrics tracking
├── sde_lib.py                  # SDE definitions
├── pure_gen_sense.py           # Pure generation (no data consistency)
├── inference_multi-coil_SENSE_h5.py  # SENSE reconstruction (15× faster)
├── inference_multi-coil_SSOS_h5.py   # SSOS reconstruction with metrics
├── batch_reconstruct_val.py    # Batch processing
├── compare_sensitivity_methods.py    # Sensitivity map tools
└── SENSE_IMPLEMENTATION.md     # Detailed SENSE implementation documentation

Technical Details

SENSE Reconstruction Pipeline

1. Sensitivity Map Estimation: Automatic estimation from k-space data using ACS region
2. Score-Based Sampling: Single combined image reconstruction with predictor-corrector sampling
3. Data Consistency: SENSE forward/adjoint operators for k-space projection
4. Efficiency: 2,000 model evaluations vs. 30,000 for per-coil method

Model Architecture

• Base Model: NCSN++ (Noise Conditional Score Network++)
• SDE: VESDE (Variance Exploding Stochastic Differential Equation)
• Training: Denoising score matching with continuous time
• Input Size: 320×320 grayscale images
• Coils: 15 channels (normalized from variable coil counts)

Performance Metrics

Method	Model Evaluations	Time (approx.)	Relative Speed
Per-Coil	30,000	~900s	1×
SENSE (Ours)	2,000	~60s	15×

Tested with N=1000 diffusion steps, M=1 corrector step on A100 GPU

System Requirements

• GPU: CUDA-capable GPU (tested on A100)
• Memory: 32GB+ RAM recommended
• Storage: ~100GB for dataset and checkpoints
• Python: 3.8+
• PyTorch: 1.10+ with CUDA support

Citation

If you use this code, please cite the original paper:

@article{chung2022score,
  title={Score-based diffusion models for accelerated MRI},
  author={Chung, Hyungjin and Ye, Jong Chul},
  journal={Medical Image Analysis},
  pages={102479},
  year={2022},
  publisher={Elsevier}
}

Related Works

• Solving Inverse Problems in Medical Imaging with Score-Based Generative Models: paper
• Robust compressed sensing using generative models: paper, code
• Original score-MRI implementation: code

Acknowledgments

This work builds upon the original score-MRI implementation by Chung and Ye. Our contributions focus on:

• HPC system integration and stability improvements
• FastMRI dataset compatibility
• SENSE-based inference optimization (15× speedup)
• Clinical multi-coil data support

License

See LICENSE file for details.