Performance comparison of medical image classification systems using TensorFlow Keras, PyTorch, and JAX

2025-07-19

AI Summary

Selecting an optimal deep learning framework for medical image analysis requires rigorous, cross-framework performance evaluation under realistic conditions. Method: This work conducts a systematic, fair benchmark of TensorFlow/Keras, PyTorch, and JAX on BloodMNIST—a blood cell classification task—using identical model architectures and training protocols across input resolutions from 28×28 to 224×224. Inference latency and classification accuracy are measured as functions of resolution. Results: JAX and PyTorch achieve state-of-the-art accuracy (>97%), significantly outperforming TensorFlow/Keras. JAX delivers the fastest inference at high resolutions (1.8× faster than PyTorch) but exhibits stronger hardware dependency. Resolution-scaling effects are nonlinear, revealing an optimal trade-off point between accuracy and efficiency. To our knowledge, this is the first empirical, multi-scale, cross-framework benchmark tailored to medical imaging, uncovering intrinsic relationships between framework design, computational efficiency, and resolution sensitivity.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Medical and Biological Imaging
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Medical imaging plays a vital role in early disease diagnosis and monitoring. Specifically, blood microscopy offers valuable insights into blood cell morphology and the detection of hematological disorders. In recent years, deep learning-based automated classification systems have demonstrated high potential in enhancing the accuracy and efficiency of blood image analysis. However, a detailed performance analysis of specific deep learning frameworks appears to be lacking. This paper compares the performance of three popular deep learning frameworks, TensorFlow with Keras, PyTorch, and JAX, in classifying blood cell images from the publicly available BloodMNIST dataset. The study primarily focuses on inference time differences, but also classification performance for different image sizes. The results reveal variations in performance across frameworks, influenced by factors such as image resolution and framework-specific optimizations. Classification accuracy for JAX and PyTorch was comparable to current benchmarks, showcasing the efficiency of these frameworks for medical image classification.

Problem

• Compares performance of TensorFlow Keras, PyTorch, and JAX for blood cell image classification
• Analyzes inference time differences and classification accuracy across frameworks
• Evaluates impact of image resolution and framework optimizations on model performance

Innovation

• Compares TensorFlow Keras, PyTorch, JAX frameworks
• Analyzes BloodMNIST dataset classification performance
• Evaluates inference time and image size impact