Technical Note: A cascade 3D U-Net for dose prediction in radiotherapy
Shuolin Liu1, Jingjing Zhang1, Teng Li1
1Key Laboratory of Intelligent Computing and Signal Processing, Ministry of Education/School of Electrical Engineering and Automation, Anhui University, HeFei, China.
Purpose:
Although large datasets are available, to learn a robust dose prediction model from a limited dataset still remains challenging. This work employed cascaded deep learning models and advanced training strategies with a limited dataset to precisely predict three-dimensional (3D) dose distribution.
Methods:
A Cascade 3D (C3D) model is developed based on the cascade mechanism and 3D U-Net network units. During model training, data augmentations are used to improve the generalization ability of the prediction model. A knowledge distillation technique is employed to further improve the capability of model learning. The C3D network was evaluated using the OpenKBP challenge dataset and competed with those models proposed by more than 40 teams globally. Additionally, it was compared with five existing cutting-edge dose prediction models. The performance of these prediction models was evaluated by voxel-based mean absolute error (MAE) and clinical-related dosimetric metrics. The code and models are publicly available online (https://github.com/LSL000UD/RTDosePrediction).
Results:
The MAE of a single C3D model without test-time augmentation is 2.50 Gy (3.57% related to prescription dose) for nonzero dose area, which outperforms the other five dose prediction models by about 0.1 Gy-1.7 Gy. The C3D model won both dose and DVH streams of AAPM 2020 OpenKBP challenge with dose score of 2.31 and DVH score of 1.55.
Conclusions:
The Cascading U-Nets is an ideal solution for 3D dose prediction from a limited dataset. The proper data preprocessing, data augmentation, and optimization procedure are more important than architectural modifications of deep learning network.
