Improving the performance of deep learning models in predicting and classifying gamma passing rates with discriminative features and a class balancing technique: a retrospective cohort study

Abstract

Background: The purpose of this study was to improve the deep learning (DL) model performance in predicting and classifying IMRT gamma passing rate (GPR) by using input features related to machine parameters and a class balancing technique.

Methods: A total of 2348 fields from 204 IMRT plans for patients with nasopharyngeal carcinoma were retrospectively collected to form a dataset. Input feature maps, including fluence, leaf gap, leaf speed of both banks, and corresponding errors, were constructed from the dynamic log files. The SHAP framework was employed to compute the impact of each feature on the model output for recursive feature elimination. A series of UNet++ based models were trained on the obtained eight feature sets with three fine-tuning methods including the standard mean squared error (MSE) loss, a re-sampling technique, and a proposed weighted MSE loss (WMSE). Differences in mean absolute error, area under the receiver operating characteristic curve (AUC), sensitivity, and specificity were compared between the different models.

Results: The models trained with feature sets including leaf speed and leaf gap features predicted GPR for failed fields more accurately than the other models (F(7, 147) = 5.378, p < 0.001). The WMSE loss had the highest accuracy in predicting GPR for failed fields among the three fine-tuning methods (F(2, 42) = 14.149, p < 0.001), while an opposite trend was observed in predicting GPR for passed fields (F(2, 730) = 9.907, p < 0.001). The WMSE_FS5 model achieved a superior AUC (0.92) and more balanced sensitivity (0.77) and specificity (0.89) compared to the other models.

Conclusions: Machine parameters can provide discriminative input features for GPR prediction in DL. The novel weighted loss function demonstrates the ability to balance the prediction and classification accuracy between the passed and failed fields. The proposed approach is able to improve the DL model performance in predicting and classifying GPR, and can potentially be integrated into the plan optimization process to generate higher deliverability plans.

Trial registration: This clinical trial was registered in the Chinese Clinical Trial Registry on March 26th, 2020 (registration number: ChiCTR2000031276). https://clinicaltrials.gov/ct2/show/ChiCTR2000031276.

Keywords: Class imbalance; Classification; Deep learning; Machine parameters; Prediction; Quality assurance.

Fig. 1

The overall workflow for the study

Fig. 2

a The proposed deep neural network based on the UNet++ architecture. b The structure of the basic convolution block

Fig. 3

Mean SHAP values computed on the validation set for each input feature in each of the eight models a–h trained during recursive feature elimination

Fig. 4

Scatter plot of measured versus predicted gamma passing rate (GPR) for models trained with the full feature set (FS1), the optimal FS (FS5), and the FS with only the Fluence map (FS8) using each of the three fine-tuning methods. The diagonal dotted line represents the perfect prediction by an ideal model. The vertical and horizontal dashed lines represent the action limit. TP = true positive; TN = true negative; FP = false positive; FN = false negative

Fig. 5

Comparison of the area under the receiver operating characteristic curve (AUC) for different models in the test dataset

Fig. 6

The distribution of SHAP values computed on the test dataset for each feature in the WMSE_FS1 (a) and WMSE_FS5 (b) models