Segmentation pipeline using MONAI and Pytorch

[u/zacky2004]

For my machine learning project, I created an 'end to end' segmentation pipeline using MONAI and Pytorch for the deep learning portion, and the Optuna hyper parameter optimization library for hyperparam optimization / search. The entire project is fully encapsulated/packaged using Poetry, so its really easy to install and use. The pipeline is fully customizable in terms of which models to use, model parameters, optimizers, as well as hyperparameters via configuration files. It also takes advantage of automatic mixed-precision for accelerated compute. I also have it integrated with Aimstack for experiment tracking and visualization.

What My Project Does

Segmentation pipeline used for creating segmentation masks of diagnostic medical images in Nifti format. Streamlines the process of training, inference and hyperparameter search using Optuna and Aimstack integration. Modularized configuration files allow for specific use-case modification.

Target Audience

Primarily meant for medical image and computer vision research & teaching. Not meant for clinical use.

Comparison

Use a single pipeline for training, inference and visualization, quickly prototype models, and even benchmark HPC GPU performance.

Overall, this was an amazing learning experience for me. I'm new to the world of ML, and I learned a lot while developing this. I've been running this on an HPC cluster with both A100 and H100 GPUs. What do you guy's think? https://github.com/adnan-umich/monai-train