Advertisement

A Multi-Task CNN Model for Autosegmentation of Prostate Patients

      Purpose/Objective(s)

      Transfer learning with deep convolutional neural networks (CNNs) has been successful in increasing the prediction accuracy of 2D classification tasks, by utilizing open source training sets such as ImageNet. Currently, there are no large public 3D segmentation data sets, so CNN models have to rely on limited training data, which can lead to poor generalization. One alternative is to create a core-shared model trained on a variety of separate but similar tasks, known as Multi-Task Learning (MTL). MTL has the ability to train on a richer and more diverse learning environment, which allows for better prediction accuracy and improved training efficiency. The present study suggests a MTL-CNN model as a superior alternative to traditional CNN models for autosegmentation of prostate patients.

      Materials/Methods

      This study evaluates the Dice coefficients between the physician-contoured ground truths (GT) and the 2D-CNN, 3D-CNN, and multi-task learning CNN (MTL-CNN) predictions. The planning CT and structures of 100 prostate patients were used in this study. Training, validation, and testing were done following Kaggle competition rules, where 80 patients were used for training, 10 patients were used for internal validation, and 10 patients were used to report final results. The 2D-CNN utilized an adaptation of U-Net, the 3D-CNN was based on a commercial 3D fully convolutional neural network architecture, and the MTL-CNN used the same core architecture as 3D-CNN but with two additional convolutional blocks per structure. The MTL-CNN was trained on all four structures simultaneously. The training time was ∼11 hours, ∼35 hours, and ∼10 hours for the 2D-CNN, 3D-CNN, and MTL-CNN respectively using a 4x data augmentation scheme for all 4 structures. All algorithms were implemented in Pytorch using two Nvidia 1080 Ti GPUs. Due to proximity and GPU memory limitations, this study was restricted to the urethra, CTV, penile bulb, and bladder. Comparator p-values were reported using Mann-Whitney U test.

      Results

      The 2D-CNN had Dice scores of 81.6 ± 2.2, 89.6 ± 1.9, 61.6 ± 1.74, and 41 ± 17.8 for the CTV, bladder, penile bulb, and urethra respectively. The 3D-CNN had Dice scores of 84.4 ± 2.15, 90.4 ± 2.1, 68.8 ± 1.3, and 60.6 ± 6.5 for the CTV, bladder, penile bulb, and urethra respectively. The MTL-CNN had Dice scores of 87.6 ± 1.9, 91.8 ± 1.94, 72.2 ± 1.72, and 68 ± 4.3. P-values of MTL-CNN versus 3D-CNN were 0.029, 0.175, 0.007, and 0.036 for the CTV, bladder, penile bulb, and urethra respectively.

      Conclusion

      This study has demonstrated that multi-task learning is capable of achieving superior training efficiency and prediction accuracy than traditional deep learning methods for CNN-based autosegmentation. Largest mean differences were seen for the most clinically challenging structures.

      Comments

      Commenting Guidelines

      To submit a comment for a journal article, please use the space above and note the following:

      • We will review submitted comments as soon as possible, striving for within two business days.
      • This forum is intended for constructive dialogue. Comments that are commercial or promotional in nature, pertain to specific medical cases, are not relevant to the article for which they have been submitted, or are otherwise inappropriate will not be posted.
      • We require that commenters identify themselves with names and affiliations.
      • Comments must be in compliance with our Terms & Conditions.
      • Comments are not peer-reviewed.