Improving Gan-Based CT Infilling Via Transfer-Learning and Layer Freezing
Abstract
Purpose
To assess the performance impact of transfer-learning and layer freezing training methods on a generative adversarial network (GAN)-based CT infilling technique.
Methods
A GAN was pretrained for 541,000 iterations using four different datasets: 9.1*106 ImageNet images, 3.7*105 publicly available head and neck CT images, 1.5*104 in-house images (NS Cancer Centre), and 3.7*105 CT images created by augmenting the in-house dataset via affine transformations. Each of the four pretrained models was fine-tuned for 50 epochs on the in-house dataset. The performance of each resulting model was scored on a validation set of in-house data and compared on each of eight metrics: precision, recall, density, coverage, Fréchet Inception Distance (FID), MSE, SSIM, and PSNR. The overall best performing pre-trained model from this transfer-learning experiment was selected for layer freezing experiments. For these experiments, nine independent sets of layers from the generator and/or critics were selected for freezing, each being used to fine-tune the best pre-trained model for 50 epochs. The metric-wise performance of each frozen fine-tuned model was compared to the best performing unfrozen fine-tuned model.
Results
The fine-tuned model pretrained using the augmented in-house data was found to outperform the other models on all eight metrics; all improvements were statistically significant (p<0.05) with respect to the model solely trained on in-house data. Freezing the shallowest half of the critics and the encoding branch of the generator provided statistically significant improvements (p<0.05) over the best performing unfrozen model on recall, FID, MSE, and SSIM. No statistically significant degradations were found using this best-tested freezing regime.
Conclusion
This work indicates that using a transfer-learning approach with an augmented dataset improves GAN-based CT infilling capability; layer freezing can further improve performance on some metrics. Future work will seek to investigate the dose distribution impacts of our technique in a radiotherapy treatment planning context.