[MRM2026] SelExNet: A Self-Supervised Physics-Informed Framework for Multi-Channel Joint RF and Gradient Waveform Optimization in 2D Spatially Selective Excitation
News | Abstract | Installation | Train | Finetune | Citation
2026.05.20 - We have released the code and pretrained weights for our paper. Please check the Installation, Train and Finetune sections for details.
2026.05.02 - Our paper is accepted by Magnetic Resonance in Medicine 2026. Work in progress.
SelExNet is a self-supervised framework for 2D spatially selective excitation that jointly optimizes radiofrequency (RF) pulses and gradient waveforms, with extension to multi-channel transmission MRI. It couples neural RF and gradient generators with a differentiable Bloch simulator, enabling pulse optimization directly from desired excitation outcomes without requiring pre-designed target pulses.
The framework designs RF pulses and parameterized variable-density spiral gradient waveforms for both single- (sTx) and multi-channel (pTx) transmission, and supports patient-specific adaptation using measured, previously unseen B0 and B1+ maps. Joint RF-Gradient optimization improves excitation fidelity over RF-only optimization. In phantom experiments, patient-specific fine-tuning restores target geometry and uniformity under field inhomogeneity. In-vivo studies demonstrate anatomically precise excitation, sharper target boundaries, and reduced off-target signal.
Figure 1: Overview of SelExNet Framework.
conda is recommended for creating the Python environment. After activating the environment, install uv and use it for faster package installation:
conda create -n selexnet python=3.10
conda activate selexnet
pip install uv
git clone https://github.com/chiew-group/SelExNet.git
cd SelExNetInstall SelExNet and its required dependencies from the repository root:
uv pip install -e .For the optional imaging dependencies, install:
uv pip install -e ".[all]"For development tools, install:
uv pip install -e ".[dev]"Core dependencies (installed via uv pip install -e .)
- torch
- omegaconf
- tqdm
- scipy
- numpy
- pytorch-msssim
- torch-pca
- opencv-python
- matplotlib
- matplotlib-inline
- Pillow
Optional imaging dependencies (installed via uv pip install -e ".[all]")
- torchvision
- pydicom
- scikit-image
- scikit-learn
Development dependencies (installed via uv pip install -e ".[dev]")
- selexnet[all]
- pytest
- black
- flake8
- isort
- mypy
- jupyterlab
python -c "import selexnet; print(selexnet.__version__)"Tip
The pretraining for 8ch pTx usually takes ~2 days, we recommend using the pretrained weights for fine-tuning, see the Finetune section.
The dataset for this project is simple which are just binary masks with different geometries. You can either download our dataset from Google Drive, and uncompress and place it in the data/train/ directory, or generate by yourself. The field maps can be in either single-map mode or multi-map mode, which is determined by the use_fm_generator flag in the configuration file for training. Format and structure of the field maps are described as follows:
-
single-map (npz format with keys:
b0,b1_real,b1_imag, unit isTesla):- b0: (H, W) array of B0 map
- b1_real: (Tx, H, W) array of real part of B1+ map
- b1_imag: (Tx, H, W) array of imaginary part of B1+ map
-
multi-map (npz format with keys:
b0_{mode},b1_real_{mode},b1_imag_{mode}where the mode can betrain,validortest, unit isTesla):- b0_train:
(N_train, H, W)array of B0 maps for training - b0_valid:
(N_valid, H, W)array of B0 maps for validation - b0_test:
(N_test, H, W)array of B0 maps for testing - b1_real_train:
(N_train, Tx, H, W)array of real part of B1+ maps for training - b1_imag_train:
(N_train, Tx, H, W)array of imaginary part of B1+ maps for training - b1_real_valid:
(N_valid, Tx, H, W)array of real part of B1+ maps for validation - b1_imag_valid:
(N_valid, Tx, H, W)array of imaginary part of B1+ maps for validation - b1_real_test:
(N_test, Tx, H, W)array of real part of B1+ maps for testing - b1_imag_test:
(N_test, Tx, H, W)array of imaginary part of B1+ maps for testing
- b0_train:
Tip
The single-map mode is for training a model that is specific to the provided field maps, while the multi-map mode is for training a model that can adapt to unseen field maps. For the multi-map mode in our experiments, we found it is not easy and stable to optimize a model that can adapt to large sets of field maps. We highly recommend using a small sets or single-map mode for training, and then fine-tuning the model on the target field maps.
The dataset should be organized in the following structure:
├──data/
│ ├── train/
│ │ ├── fm/
│ │ │ ├── field_map.npz # contains B0 and B1+ maps
│ │ │ └── loading_mask.npy # contains the binary mask of the loading
│ │ └── dataset/
│ │ ├── sample_1.png # sample image 1
│ │ ├── sample_2.png # sample image 2
│ │ └── ... # more samples
│ ├── finetune/
python main.py \
--cfg configs/train_config.yaml \
--resume # resume the training from the last checkpoint, otherwise ignore this flagIf you have multiple GPUs, you can use distributed training to speed up the training process. You can launch it with the following command:
torchrun --nproc_per_node=<num_gpus> --standalone main.py \
--cfg configs/train_config.yaml \
--resume \ # resume the training from the last checkpoint, otherwise ignore this flag
--ddp # use distributed data parallel trainingThe configuration file is in YAML format. You can find an example configuration file at configs/train_config.yaml. You can also modify the parameters in the configuration file according to your needs. All the training outputs, including the trained model checkpoints and the training logs, will be saved in the outputs/exp_name directory.
After training the model, you can fine-tune it on the unseen field maps to further improve the excitation performance. We provided pretrained weights, example field maps and shape which can be used for fine-tuning. You can download them:
-
Pretrained weights: place it in the
data/finetune/pretrained/directory. -
Example field maps: place it in the
data/finetune/fm/directory. -
Example phantom mask: place it in the
data/finetune/fm/directory.
Important
Make a copy of excitation shape image to the data/finetune/dataset/ directory for fine-tuning (see the example in data/finetune/dataset/). This is due to the fact that the fine-tuning is a kind of overfitting process which needs same image for both training and validation.
python finetune.py \
--cfg configs/finetune_config.yaml \
--finetune_weights <path/to/finetune/pretrained.pt>For demo, the --finetune_weights should be set to data/finetune/pretrained/pretrained.pt. The configuration file is in YAML format. You can find an example configuration file at configs/finetune_config.yaml. You can also modify the parameters in the configuration file according to your needs. All the fine-tuning outputs, including the fine-tuned model checkpoints and the fine-tuning logs, will be saved in the outputs/exp_name directory. Once the fine-tuning is done, you can use the fine-tuned model for inference on the target excitation shape and field maps.
python test_single.py \
--cfg configs/finetune_config.yaml \
--output_dir <path/to/output/dir> \
--img_path <path/to/shape/image.png>In the demo, the --output_dir is set to demo/ and --img_path should be data/finetune/dataset/shape.png. Finally, the simulated excitation pattern, and tailored RF and gradient waveforms will be saved in demo/ directory.
Note
The designed RF and gradient waveforms will be saved in .mat format which is compatible with the native Siemens pTx workflow. If you need to use on other platforms, you should manually convert them to the required format.
If you find this code useful for your research, please consider citing our paper:
@article{xiao2026selexnet,
title={SelExNet: A Self-Supervised Physics-Informed Framework for Multi-Channel Joint RF and Gradient Waveform Optimization in 2D Spatially Selective Excitation},
author={Yuliang Xiao and Jason Rock and Zhe Wu and Jamie Near and Mark Chiew and Simon J. Graham},
journal={Magnetic Resonance in Medicine},
year={2026},
publisher={Wiley Online Library}
}