Training of a MEDNEXT model

[plbenveniste]

This issue describes the work done to train a MedNext model and to evaluate its performance.

Work is done under the branch plb/monai_unet.

The code containing the MedNext model is in this repo: https://github.com/MIC-DKFZ/MedNeXt

[plbenveniste]

Installation and setup

conda create -n venv_mednext python=3.9
conda activate venv_mednext
pip install -r ms-lesion-agnostic/monai/requirements.txt
git clone https://github.com/MIC-DKFZ/MedNeXt
cd MedNext
pip install -e .

Training

CUDA_VISIBLE_DEVICES=1 python ms-lesion-agnostic/monai/train_monai_mednext_lightning.py --config ms-lesion-agnostic/monai/config.yml

[plbenveniste]

I am also training (on kronos) two other MedNext model : input_channel = 32 or 64. To have sufficient memory for this, I reduced the number of samples produced by randCropByPosNeg from 4 to 2 and 1.

[plbenveniste]

I am currently running inference using the code in monai/test_model_mednext.py

The code was ran with the following command:

CUDA_VISIBLE_DEVICES=1 python ms-lesion-agnostic/monai/test_model_mednext.py --config ms-lesion-agnostic/monai/config_test.yml --data-split test

The plots were computed with the following command:

python ms-lesion-agnostic/monai/plot_performance.py --pred-dir-path /home/plbenveniste/net/ms-lesion-agnostic/results/2024-09-02_12:14:28.124188/test_set/ --data-json-path ~/net/ms-lesion-agnostic/msd_data/dataset_2024-07-24_seed42_lesionOnly.json --split test

Here are the results:

[plbenveniste]

In a new experiment, I am training a MedNext model using the nnUNet framework:

export nnUNet_raw_data_base="/home/plbenveniste/net/ms-lesion-agnostic/mednext_experiment/nnUNet_raw_data_base"
export nnUNet_preprocessed="/home/plbenveniste/net/ms-lesion-agnostic/mednext_experiment/mednext_preprocessed"
export RESULTS_FOLDER="/home/plbenveniste/net/ms-lesion-agnostic/mednext_experiment/mednext_results"

Because it uses nnUNetv1, I had to do the following changes:

• Changed 'Dataset301_msLesionAgnostic' to 'Task301_msLesionAgnostic'
• Changed structure from 'nnUNet_raw/Dataset301_msLesionAgnostic' to 'nnUNet_raw_data_base/nnUNet_raw/Task301_msLesionAgnostic'
• Changed 'num_channels' to 'modality'
• Changed labels from "background": 0 to "0":"background"

I also installed nnunet v1:

git clone https://github.com/MIC-DKFZ/nnUNet.git
cd nnUNet
git checkout nnunetv1
pip install -e .

To preprocess:

mednextv1_plan_and_preprocess -t YOUR_TASK -pl3d ExperimentPlanner3D_v21_customTargetSpacing_1x1x1

Then for training, I chose a 3x3 kernel which is pretty deep (large version: block_counts = [3,4,8,8,8,8,8,4,3])

mednextv1_train 3d_fullres nnUNetTrainerV2_MedNeXt_L_kernel3 301 0 -p nnUNetPlansv2.1_trgSp_1x1x1

The results of this experience were inconclusive !