Model Overview
AttentionUNet for MRI Segmentation on the SKMTEA dataset.
ATOMMIC: Training
To train, fine-tune, or test the model you will need to install ATOMMIC. We recommend you install it after you've installed latest Pytorch version.
pip install atommic['all']
How to Use this Model
The model is available for use in ATOMMIC, and can be used as a pre-trained checkpoint for inference or for fine-tuning on another dataset.
Corresponding configuration YAML files can be found here.
Automatically instantiate the model
pretrained: true
checkpoint: https://huggingface.co/wdika/SEG_DynUNet_SKMTEA/blob/main/SEG_DynUNet_SKMTEA.atommic
mode: test
Usage
You need to download the SKM-TEA dataset to effectively use this model. Check the SKMTEA page for more information.
Model Architecture
model:
model_name: SEGMENTATIONDYNUNET
segmentation_module: DYNUNet
segmentation_module_input_channels: 1
segmentation_module_output_channels: 4
segmentation_module_channels:
- 32
- 64
- 128
- 256
- 512
segmentation_module_kernel_size:
- 3
- 3
- 3
- 3
- 1
segmentation_module_strides:
- 1
- 1
- 1
- 1
- 1
segmentation_module_dropout: 0.0
segmentation_module_norm: instance
segmentation_module_activation: leakyrelu
segmentation_module_deep_supervision: true
segmentation_module_deep_supervision_levels: 2
segmentation_module_normalize: false
segmentation_module_norm_groups: 2
segmentation_loss:
dice: 1.0
dice_loss_include_background: true # always set to true if the background is removed
dice_loss_to_onehot_y: false
dice_loss_sigmoid: false
dice_loss_softmax: false
dice_loss_other_act: none
dice_loss_squared_pred: false
dice_loss_jaccard: false
dice_loss_flatten: false
dice_loss_reduction: mean_batch
dice_loss_smooth_nr: 1e-5
dice_loss_smooth_dr: 1e-5
dice_loss_batch: true
dice_metric_include_background: true # always set to true if the background is removed
dice_metric_to_onehot_y: false
dice_metric_sigmoid: false
dice_metric_softmax: false
dice_metric_other_act: none
dice_metric_squared_pred: false
dice_metric_jaccard: false
dice_metric_flatten: false
dice_metric_reduction: mean_batch
dice_metric_smooth_nr: 1e-5
dice_metric_smooth_dr: 1e-5
dice_metric_batch: true
segmentation_classes_thresholds: [0.5, 0.5, 0.5, 0.5]
segmentation_activation: sigmoid
magnitude_input: true
log_multiple_modalities: false # log all modalities in the same image, e.g. T1, T2, T1ce, FLAIR will be concatenated
normalization_type: minmax
normalize_segmentation_output: true
complex_data: false
Training
optim:
name: adam
lr: 1e-4
betas:
- 0.9
- 0.98
weight_decay: 0.0
sched:
name: InverseSquareRootAnnealing
min_lr: 0.0
last_epoch: -1
warmup_ratio: 0.1
trainer:
strategy: ddp_find_unused_parameters_false
accelerator: gpu
devices: 1
num_nodes: 1
max_epochs: 20
precision: 16-mixed # '16-mixed', 'bf16-mixed', '32-true', '64-true', '64', '32', '16', 'bf16'
enable_checkpointing: false
logger: false
log_every_n_steps: 50
check_val_every_n_epoch: -1
max_steps: -1
Performance
Evaluation can be performed using the segmentation evaluation script for the segmentation task, with --evaluation_type per_slice.
Results
Evaluation
DICE = 0.6888 +/- 0.1359 F1 = 0.05911 +/- 0.2638 HD95 = 8.973 +/- 4.507 IOU = 0.01517 +/- 0.06638
References
[1] ATOMMIC
[2] Desai AD, Schmidt AM, Rubin EB, et al. SKM-TEA: A Dataset for Accelerated MRI Reconstruction with Dense Image Labels for Quantitative Clinical Evaluation. 2022