TLDR: This paper compares TensorFlow Keras, PyTorch, and JAX for classifying blood cell images from the BloodMNIST dataset. It evaluates their inference times and classification accuracy for 28×28 and 64×64 image resolutions. PyTorch showed the fastest inference for smaller images, while JAX excelled with larger images. All frameworks achieved comparable high classification accuracy, with minor variations depending on image size. The study highlights the importance of framework choice based on specific task requirements and image characteristics.
Medical imaging is a cornerstone of modern healthcare, playing a crucial role in the early detection and ongoing monitoring of various diseases. Among its many applications, blood microscopy stands out for its ability to offer invaluable insights into the intricate morphology of blood cells, aiding significantly in the diagnosis of hematological disorders and infections. In recent years, the advent of deep learning has revolutionized this field, leading to automated classification systems that promise to enhance both the accuracy and efficiency of blood image analysis, becoming indispensable tools for pathologists and clinicians alike.
Despite the widespread adoption of deep learning in medical imaging, a persistent challenge remains: selecting the most efficient framework for deployment. This decision is complicated by variations in execution speed, optimization capabilities, and support for specialized hardware accelerators. A recent research paper, titled “Performance comparison of medical image classification systems using TensorFlow Keras, PyTorch, and JAX” by Merjem Bećirović, Amina Kurtović, Nordin Smajlović, Medina Kapo, and Amila Akagić, directly addresses this challenge. The study meticulously compares three prominent deep learning frameworks – TensorFlow with Keras, PyTorch, and JAX – under identical model and training configurations. Their primary goal was to evaluate inference speed and classification accuracy on blood cell microscopy images, thereby identifying which framework offers the most efficient execution for medical image classification tasks. Understanding these trade-offs is vital for researchers and practitioners aiming to deploy robust and scalable systems in clinical environments where both precision and speed are paramount.
The Deep Learning Frameworks in Focus
The paper delves into the characteristics of each framework:
- TensorFlow and Keras: Released by Google Brain in 2015, TensorFlow is renowned for its scalability and performance optimizations, representing computations as directed graphs. Keras, its high-level API, simplifies model building. TensorFlow supports both static and dynamic graph execution and can integrate with XLA for optimized linear algebra.
- PyTorch: Introduced by Meta AI in 2016, PyTorch balances usability and speed with an imperative, Pythonic programming style. It pioneered dynamic eager execution, which improved flexibility and debugging. PyTorch also offers JIT (Just-In-Time) compilation and XLA for optimization, though they need to be explicitly enabled.
- JAX: A 2018 Google research project, JAX focuses on accelerator-driven array computation and program transformation. It leverages the XLA compiler to optimize code for various hardware architectures. A core feature of JAX is its JIT compilation, which transforms Python functions into efficient, accelerable code by tracing their execution.
Methodology and Dataset
The researchers utilized the BloodMNIST dataset, a publicly available collection of 17,092 2D blood cell microscopy images from healthy individuals, part of the larger MedMNIST collection. The dataset is pre-split into training, validation, and testing subsets and contains images at various resolutions. For this study, due to computational constraints, the focus was on images resized to 28×28 and 64×64 pixels. The task involved classifying images into 8 distinct blood cell types.
A custom ResNet-inspired convolutional neural network (CNN) was designed for the experiments, balancing computational efficiency and performance. The model featured an initial convolutional layer, followed by six residual blocks with increasing filter numbers, and concluded with a global max-pooling layer and a fully connected layer for classification. The model was trained for 20 epochs using the Adam optimizer and categorical cross-entropy loss.
Key Findings: Inference Time and Classification Performance
The study’s results revealed distinct performance variations across the frameworks:
- Inference Times: For the smaller 28×28 images, PyTorch demonstrated the fastest average inference time, significantly outperforming TensorFlow with Keras and JAX. This suggests PyTorch’s superior efficiency for smaller image sizes, possibly due to its built-in optimization strategies handling simple operations more efficiently than JAX’s JIT compilation overhead for such small tasks. However, for the larger 64×64 images, JAX took the lead with the fastest inference time, surpassing PyTorch, while Keras showed the slowest performance. This highlights JAX’s strength in handling larger images efficiently, where its JIT compilation benefits become more pronounced.
- Classification Performance: Despite identical architectures and settings, the frameworks showed interesting variations in classification accuracy. For 28×28 images, JAX consistently held an edge over both TensorFlow with Keras and PyTorch across all evaluation metrics (accuracy, precision, recall, F1-score). For 64×64 images, PyTorch achieved the highest accuracy (98.22%), closely followed by JAX (97.9%) and TensorFlow with Keras (97.63%). PyTorch also led in macro and weighted averages for precision, recall, and F1-score for this image size. These differences may stem from framework-specific implementations, such as optimization strategies or computational precision, influencing model training dynamics.
Also Read:
- AI Model for Anemia Detection: Balancing Performance and Efficiency on Portable Devices
- RARE-UNet: A Smart Approach to Medical Image Segmentation Across Different Resolutions
Conclusion and Future Outlook
The research concludes that while inference time is a critical factor in framework selection, the differences in classification performance can be attributed to nuanced factors like weight initialization, optimization algorithms, and JIT compilation overhead. PyTorch proved fastest for smaller images, and JAX for larger ones, while both demonstrated comparable, high classification accuracy. The study acknowledges a limitation in using relatively low-resolution images (28×28 and 64×64 pixels) compared to typical clinical digital microscopy. However, the authors suggest that these findings offer valuable insights into framework behavior and scaling trends that are likely to generalize to higher resolutions. Future work could involve integrating super-resolution methods to enhance image quality and dimensions, making the study’s applicability even broader. For more detailed information, you can read the full research paper here.
