clara_pt_net_arch_search_segmentation

Model Overview

A pre-trained model for volumetric (3D) segmentation of the pancreas and tumour from CT image.

Note: The 4.1 version of this model is only compatible with the 4.1 version of the Clara Train SDK container

Model Architecture

This model is trained using the neural network model from the neural architecture search algorithm, DiNTS [1].

Training

Neural Architecture Search Configuration

The neural architecture search was performed with the following:

• Script: search_multi_gpu_byow.sh
• GPU: At least eight GPUs with 16GB of GPU memory
• Actual Model Input: 96 x 96 x 96 for traing, 96 x 96 x 96 for validation/testing
• AMP: True
• Optimizer: SGD
• Initial Learning Rate: 0.025
• Loss: DiceCELoss

Training Configuration

The training was performed with the following:

• Script: train_multi_gpu_byow.sh
• GPU: (at least) eight 16GB of GPU memory
• Actual Model Input: 96 x 96 x 96 for traing, 96 x 96 x 96 for validation/testing
• AMP: True
• Optimizer: SGD
• (Initial) Learning Rate: 0.025
• Loss: DiceCELoss

Note: If out-of-memory or program crash occurs while caching the data set, please change the cache_rate in CacheDataset to a lower value in the range (0, 1).

The segmentation of pancreas region is formulated as the voxel-wise 3-class classification. Each voxel is predicted as either foreground (pancreas body, tumour) or background. And the model is optimized with gradient descent method minimizing soft dice loss [2] and cross-entropy loss between the predicted mask and ground truth segmentation.

Dataset

The training data is from the Medical Decathlon.

• Target: Pancreas and tumor
• Task: Segmentation
• Modality: CT
• Size: 281 3D volumes (196 Training, 56 Validation, and 29 Testing)
• Challenge: Large ranging foreground size

Performance

Dice score is used for evaluating the performance of the model. The trained model achieved average testing Dice score 0.5599 over 29 volumes (0.7661 for class 1, and 0.3538 for class 2).

Training

Training loss over 2920 epochs.

Validation

Validation mean dice score over 2920 epochs.

Searched Architecture Visualization

Users can install Graphviz for visualization of searched architectures (needed in custom/decode_plot.py). The edges between nodes indicate global structure, and numbers next to edges represent different operations in the cell searching space. An example of searched architecture is shown as follows:

How to Use this Model

The model was validated with NVIDIA hardware and software. For hardware, the model can run on any NVIDIA GPU with memory greater than 16 GB. For software, this model is usable only as part of Transfer Learning & Annotation Tools in Clara Train SDK container. Find out more about Clara Train at the Clara Train Collections on NGC.

Full instructions for the training and validation workflow can be found in our documentation.

Input

Input: 1 channel CT image with intensity in HU

1. Converting to channel first
2. Normalizing and clipping intensities of tissue window to [0,1]
3. Cropping foreground surrounding regions
4. Cropping random fixed sized regions of size [96, 96, 96] with the center being a foreground or background voxel at ratio 1 : 1
5. Randomly rotating volumes
6. Randomly zooming volumes
7. Randomly smoothing volumes with Gaussian kernels
8. Randomly scaling intensity of the volume
9. Randomly shifting intensity of the volume
10. Randomly adding Gaussian noises
11. Randomly flipping volumes

Output

Output: 3 channels

• Label 0: background
• Label 1: pancreas body
• Label 2: pancreas tumour

Sliding-window Inference

Inference is performed in a sliding window manner with a specified stride.

Limitations

This training and inference pipeline was developed by NVIDIA. It is based on a segmentation model developed by NVIDIA researchers. This research use only software has not been cleared or approved by FDA or any regulatory agency. Clara pre-trained models are for developmental purposes only and cannot be used directly for clinical procedures.

References

[1] He, Yufan, et al. "Dints: Differentiable neural network topology search for 3d medical image segmentation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021. https://openaccess.thecvf.com/content/CVPR2021/papers/He_DiNTS_Differentiable_Neural_Network_Topology_Search_for_3D_Medical_Image_CVPR_2021_paper.pdf

[2] Milletari, Fausto, Nassir Navab, and Seyed-Ahmad Ahmadi. "V-net: Fully convolutional neural networks for volumetric medical image segmentation." 2016 fourth international conference on 3D vision (3DV). IEEE, 2016. https://arxiv.org/abs/1606.04797.

License

End User License Agreement is included with the product. Licenses are also available along with the model application zip file. By pulling and using the Clara Train SDK container and downloading models, you accept the terms and conditions of these licenses.