Runtime of UNet Breast Cancer Segmenter

[rictoo]

• TIA Toolbox version: 1.5.1
• Python version: 3.9.18
• Operating System: CentOS

Description

Hi all! I've been trying to perform semantic segmentation on a ~60k x 60k pixel WSI (20x magnification) of breast tissue using the pre-trained 'fcn_resnet50_unet-bcss' model. I'm finding that it performs very well on my images, but that it is taking a very long time to run inference. On a V100, with a batch size of 32 and 8 CPU cores, it takes about 30-50 minutes to run on a single WSI. I would ideally prefer that it run much faster than this, as I plan to deploy the model on thousands of images.

What I Did

I've tried running the model with various 'input_resolutions'/'output_resolutions' (from the default 0.25 to 1, and 4), however this resulted in much poorer performance.

My set-up looks like this:

bcc_segmentor = SemanticSegmentor(
    pretrained_model="fcn_resnet50_unet-bcss",
    num_loader_workers=8,
    batch_size=32,
    auto_generate_mask=False,
)

bcc_wsi_ioconfig = IOSegmentorConfig(
    input_resolutions=[{"units": "mpp", "resolution": 0.25}],
    output_resolutions=[{"units": "mpp", "resolution": 0.25}],
    patch_input_shape=[1024, 1024],
    patch_output_shape=[512, 512],
    stride_shape=[512, 512],
    save_resolution={"units": "mpp", "resolution": 2},
)

# WSI prediction
wsi_output = bcc_segmentor.predict(
    imgs=[wsi_file_name],
    masks=None,
    save_dir=save_dir,
    mode="wsi",
    ioconfig=bcc_wsi_ioconfig,
    on_gpu=ON_GPU,
    crash_on_exception=True,
)

I'm wondering whether you have any suggestions on the ideal way to decrease runtime? I'm happy to sacrifice a bit of quality. However, I have a feeling that since the model was trained on a specific resolution, decreasing the input resolution might result in unacceptably extremely poor performance. Was I on the right track by changing "input_resolutions" and "output_resolutions"? Or would I need to maybe retrain the entire model at a lower resolution (using the BCSS dataset)?

Thanks for the amazing project!

[measty]

Hi @rictoo , One thing you could try if your slides have a fair amount of white space is to provide a mask, as otherwise the model will be needlessly processing all non-tissue regions of the slide. This can be done either by setting 'auto_generate_mask' to True (in which case tiatoolbox will generate a tissue mask on the fly for each slide) or pass a mask explicitly to the predict method if you have a tissue mask from some external source (for example HistoQC) you'd rather use. The segmentation will then be done only in the tissue regions, which could cut the run time by a significant chunk.

[rictoo]

Hi @measty, thanks for your response. Indeed, I have noticed a significant speed improvement since using a mask. However, the speed at which it performs inference on a slide with the otherwise default settings was still very slow, unfortunately.

Since making this post, I actually ended up fine-tuning the fcn_resnet50_unet-bcss model on 1mpp image tiles sourced from the original training set which can be found here. This resulted in significant performance improvements and runs at 16x the original speed (due to being 4x smaller in both height and width). I would, however, be curious to know what loss function was used in the original training of the model. Was it Dice? BCE? Something else?

[shaneahmed]

@simongraham Please can you confirm what loss function was used to train these models?

[mostafajahanifar]

Hi @rictoo,

If you are sure that you are using ON_GPU=True, another thing that you can try is to set the strid to be equal to the patch size, that should also give your 2x boost in speed, however may degrade the quality on patch boundaries.

About the loss function, I think this model was trained with BCE only (if I'm not wrong! @vqdang knows better as he trained this model). However, if you are thinking of fine-tuning it yourself, maybe a combination of Dice+BCE would be a better idea.

Good luck.