MedMNIST (PathMNIST) Active Learning

DeterminedAI WebUI
Github ReadMe
Model on Hugging Face

Inspiration

Medical imaging often requires vast amounts of labeled data, which can be expensive and time-consuming to obtain. Active learning offers a way to minimize this burden by identifying the most informative data for labeling. This project applies active learning to the MedMNIST PathMNIST dataset, aiming to improve classification accuracy with fewer labeled examples. The goal is to explore scalable, efficient solutions for building medical diagnostic tools.

What it does

The project uses active learning to train a ResNet-50 model for classifying pathological tissue types. It combines uncertainty sampling and clustering to select the most informative samples for labeling. The model learns iteratively, improving its predictions as additional samples are labeled and included in the training set.

How we built it

The project was built using Determined AI’s platform for managing machine learning workflows. We used the PathMNIST dataset from MedMNIST and applied data augmentation techniques to improve model robustness. The ResNet-50 architecture was customized to handle 28x28 pixel images and predict nine classes.

Training was conducted in an active learning loop. Each iteration included training the model, evaluating its performance, and selecting new samples for labeling using a combination of uncertainty scores and k-means clustering. Determined AI's platform streamlined the process by providing tools for managing experiments and monitoring metrics.

Challenges we ran into

Balancing uncertainty and diversity in the active learning strategy was challenging. Selecting samples that were both uncertain and representative of the dataset required careful tuning of the sampling algorithm. Additionally, adapting ResNet-50 to smaller inputs while maintaining performance involved modifying its layers and optimizing training parameters.

In addition to this, Determined AI was run on a laptop that is used for school and working with all of the resources being pooled into Determined AI was difficult. Initially, the model was developed on a local computer using a HPC to test for errors, but as time went on and I learned more about Determined AI's services, I switched to Determined AI. However, converting the code was a pretty big challenge as I have never experienced it before.

Also, there were some challenges to upload the model to Hugging Face as the model checkpoints were saves as .pt and not as as model.bin file and this complicated the upload process.

Accomplishments that we're proud of

We achieved a model accuracy of 94.72% on the PathMNIST dataset, surpassing the MedMNIST benchmark accuracy of 91.1%. This demonstrates the potential of active learning to enhance performance with limited labeled data. Successfully integrating Determined AI's platform into the workflow was another significant accomplishment, enabling efficient training and evaluation.

What we learned

We gained a deeper understanding of active learning and its practical applications in reducing labeling effort. The project also highlighted the importance of data diversity in improving model generalization. Working with Determined AI provided valuable insights into managing iterative training workflows at scale.

What's next for MedMNIST Active Learning

Next steps include extending this approach to other datasets in the MedMNIST collection and refining the active learning strategies. We also plan to explore how transfer learning can generalize the model to new medical imaging tasks and scale this approach for real-world diagnostic applications.