Improving OCT B-scan of interest inference performance using TensorRT based neural network optimization

Abstract

Purpose : OCT B-scan of interest detection is a newly developed deep learning algorithm that aims to improve the workflow efficiency for CIRRUS™ HD-OCT during OCT data review. However, the inference time of the neural network can be long for 3D scans that have a large number of B-scans. In this study, we propose a method that can improve the inference performance of OCT B-scan of interest algorithm by optimizing the neural network using TensorRT.

Methods : 76,544 OCT B-scans from 598 macular cubes acquired from 598 subjects using CIRRUS HD-OCT 5000 (ZEISS, Dublin, CA) were used as the training and development sets. 148 B-scans were ungradable and were excluded. A 3-channel ResNet-50 neural network was trained using 61,058 B-scans from the training set and fine-tuned using 15,338 B-scans from the development set. The trained neural network was then frozen and saved as a protocol buffers (pb) file.
The frozen model was loaded into memory using TensorFlow 1.13.1 and optimized with TensorRT 5.0.2. A maximum batch size of 128 was specified when the TensorRT engine optimized the neural network.
To test the speed difference between the frozen model and the optimized model, 12,800 inferences were performed and the average inference time for each B-scan was measured. To test the accuracy of the optimized model, 25,600 B-scans acquired from 200 subjects at 3 different clinical sites were used. Both the B-scan of interest binary predication result and probability were compared between the original frozen model and the optimized model.

Results : For model complexity, the number of nodes of the network was reduced from 881 to 3 after optimization. For inference speed, each B-scan inference took 12.6 milliseconds before optimization and 6.9 milliseconds after optimization. The inference time for a macular cube scan was 1.61 seconds before optimization and 0.88 second after optimization. The optimized model ran about 1.8 times the speed of the original model. For accuracy, the sensitivity, specificity and AUC on the test set before and after optimization were the same: 92.6%, 95.1% and 0.98. The mean probability difference between the frozen model and the optimized model was negligible (-4.3e-8 ±5.0e-7).

Conclusions : In this study, we demonstrated a method that improves the inference speed for OCT B-scan of interest detection algorithm with the same accuracy maintained.

This is a 2020 ARVO Annual Meeting abstract.