vanvalenlab / deepcell-tf

Deep Learning Library for Single Cell Analysis
https://deepcell.readthedocs.io
Other
427 stars 99 forks source link

Processing large Whole Slide Images with MESMER- Potential memory issues? #582

Open rjesud opened 2 years ago

rjesud commented 2 years ago

Hi, Thanks for developing this useful tool! I’m eager to apply it in our research. I am encountering what I believe is an out of memory crash when processing a large digital pathology whole slide image (WSI). This specific image Is 76286 x 44944 (XY) at .325 microns/pixel. While this is large, it is not atypical in our WSI datasets. Other images in our dataset can be larger. I have done benchmarking with the data and workflow described in another ticket here: https://github.com/vanvalenlab/deepcell-tf/issues/553. I can confirm results and processing time (6min) as shown here: https://github.com/vanvalenlab/deepcell-tf/issues/553#issuecomment-940622124. So it appears things are working as expected. So, I am not sure where I am hitting an issue with my image. I assumed that the image was being sent to GPU in memory-efficient batches. Perhaps this is occurring at the pre-processing or post-processing stage? Should I explore a different workflow for Mesmer usage with images of this size?

I am using: Deep-cell 0.11.0 Tensorflow 2.5.1 Cudnn 8.2.1 Cudotoolkit 11.3.1 GPU: Quadro P6000 computeCapability: 6.1 coreClock: 1.645GHz coreCount: 30 deviceMemorySize: 23.88GiB deviceMemoryBandwidth: 403.49GiB/s

CPU: 16 cores, 128 GB RAM

Seems to be talking to GPU: image

image

Unfortunately, does not complete: image

msschwartz21 commented 2 years ago

Hi rjesud, I expect @ngreenwald may have some specific advice to add, but in the meantime here's a few things to look at.

One notable difference between your data and the benchmarks described in #553 is that you specify image_mpp=0.325. The model was trained with 0.5, so the first thing the application does is to resize your image to match the resolution that the model was trained with. If you can, please try the options described below and let us know the results so that we can diagnose the problem.

  1. Run predictions on your data without specifying the resize. The results may not be perfect, but it will confirm that the rest of the system is working correctly. labeled_image = app.predict(im, compartment='whole-cell')
  2. Run the resize function in isolation and see if it completes in a reasonable time:

    from deepcell_toolbox.utils import resize
    
    shape = im.shape
    image_mpp = 0.325
    model_mpp = app.model_mpp
    
    scale_factor = image_mpp / model_mpp
    new_shape = (int(shape[1] * scale_factor),
                         int(shape[2] * scale_factor))
    image = resize(image, new_shape, data_format='channels_last')
ngreenwald commented 2 years ago

There's indeed a post-processing step that operates across the entire image. I'm guessing for something this large, it will make more sense to generate segmentation predictions on large patches (say 10k x 10k pixels), and then stitch them together into a single image. This isn't something that's currently supported by the application itself, but you could implement it with a for loop. If it ends up being useful, we could discuss incorporating it as an option inside the application, perhaps with some intelligent behavior at the border between patches.

rjesud commented 2 years ago

@msschwartz21,

Thank you for the suggestion.

  1. Running the prediction without specifying resize did not complete. Again, I assume out of memory.
  2. Interestingly, the resize function also fails.

@ngreenwald, tile/patch-wise post processing would be great! Can we add this as a feature request?

ngreenwald commented 2 years ago

Absolutely, feel free to open a separate detailing the requirements and what an ideal interface would look like. No promises on having something immediately ready. In the meantime, something like this should work for you:

labeled_image = np.zeros_like(im)
step_size = 10000

for row in range(0, labeled_image.shape[0], step_size)
      for col in range(0, labeled_image.shape[1], step_size)
           labeled_image[row:(row + step_size), col:(col + step_size)] = / 
                     app.predict(im[row:(row + step_size), col:(col + step_size)]
rjesud commented 2 years ago

@ngreenwald Thanks for this snippet! I'll give it a try. We'll probably have to get fancy with tile overlaps to solve the border cell issue.

ngreenwald commented 2 years ago

Yes, although with 10k x 10k crops, the number of cells at the border would quite a small overall percentage of the image, might be fine to just leave them as is. Up to you of course.

rjesud commented 2 years ago

Hi, I have a follow up question. Is the predict function executing normalization of the image array prior to inference? If so, is it possible to override this step and perform it outside of the function. This would be the preferred method to ensure each tile is using the same normalization scaling.

ngreenwald commented 2 years ago

Yes, you can pass kwarg to the function to control whether normalization happens or not, see here. However, I would verify that you see a decrease in performance before substituting your own normalization. The normalization we're using, CLAHE, already adjusts across the image to some degree, so I don't know that it will necessarily cause issues to have it applied to different tiles, especially tiles as large as 10k x 10k.

The only area where I've seen issues before is if entire tiles are devoid of signal, and it's just normalizing background to background.