Saving predictions occasionally during training and saving best predictions during validation using multiple validator dataloaders.

[ndahiya3]

I am working on 2D/3D medical image segmentation tasks and I have been trying to use lightning to simplify/reduce my boiler plate code. All of this code is functional in simple pytorch code but is large, messy and difficult to generalize from project to project, hence using lightining.

One of the things I am stuck on right now is to be able to occasionally save training batches with predictions (2D/3D float images). I can return the predictions from training_step but I am not sure about the memory footprint of doing that? There are several 100k's 512x512 images and obviously can't store all of them in memory. So, is there any built-in solution that I am missing? Any pointers are appreciated.

Second major thing I am stuck on is using multiple validation dataloaders in evaluation step. There are several things here. Each of the validation dataloaders combines 2D image slices belonging to a single 3D dicom image. Using several dataloaders I can pass in several 3D images separately to a 2D model. After predicting segmentation maps on all slices from a single dataloader, I need to compute some metrics (dice, hausdorff etc.) taking into account all the prediction but only from single dataloader. Any pointers on how to do that? To compute 3D hausdorff distance I have to keep all the predictions from a single dataloader in memory. There's no way around it. But I can't store all the images of all the dataloaders at once.

Continuing the above point. Once I have computed metrics from each dataloader, I need to average the metrics and if average has improved compared to the previous epoch, I need to save the predictions/outputs as separate 3D images, with each 3D image belonging to one dataloader. Again, I am not sure how to do that without just saving (in RAM) the predictions from all the validators together. This will again have a huge memory footprint.

So far, my thought is to figure out if there is a signal or a hook which indicates one dataloader is finished. If I know that, then I can compute the metrics for this dataloader, save the predictions as 3D and then CLEAR the RAM storage. This can then be repeated for each dataloader. Each dataloader only has a few hundred slices, so I can store all those predictions in RAM at a time. At the end of evaluation I could manually compute average of metrics, compare to saved average from earlier epochs, and if there is an improvement then I can just rename the saved 3D images to indicate improved validation results.

Any thoughts on these issues are appreciated. I have already spent a bunch of time trying figure this out.