Speed up prediction by chunks

[danifranco]

We need to rethink how the final prediction is reconstructed using inference by chunks, as right now we are using so much disk (and time) with the mask that needs to be divided to take into account the overlap/padding as here.

[ClementCaporal]

Hi @danifranco, I hope you are doing well, I will look at this issue next week and send an associated PR!

[danifranco]

Hi @ClementCaporal, glad to hear from you! I was planning to change how the inference is done too and speed up the process by incorporating a Pytorch generator. What do you think if I first modify that and then you try to rethink this disk related problem?

[ClementCaporal]

It depends if the inference with the pytorch generator really change this cropping part (I suppose not?)? If yes, I better wait!

BiaPy/biapy/engine/base_workflow.py

Line 1127 in f872e1f

p = p[0, self.cfg.DATA.TEST.PADDING[0]:p.shape[1]-self.cfg.DATA.TEST.PADDING[0],

And for the inference structure, how do you think it will compare with: https://thejacksonlaboratory.github.io/zarrdataset/examples/advanced_example_pytorch_inference.html ?

[danifranco]

Yeah I think that will be changed too. I have an idea in mind: let me cook for two weeks or so and see if I can have a first version of the generator. Then we both can build upon, correcting bugs that may appear and including the zarrdataset to better handle Zarrs (not sure about this last as I would like to maintain H5 large file support too). What do you think?

[ClementCaporal]

Can't wait to taste what you've cooked! See you in two weeks

+1 to not drop .h5 support

[danifranco]

TLDR:
This new way of doing the "by chunks" inference should be faster than the one we had plus it does not require a mask to make the division of the final Zarr.

Long explanation:
I have just finished the changes!! I completely redo the by chunks inference and now it works like this:

• For each test sample:
  • A data loader is created (chunked_test_pair_data_generator) which is in charge of sharding the Zarr/H5 data. It takes into account the number of workers + the number of GPUs. The usage of the dataloader is key point comparing to our old implementation as it should speed up the process because of the underlying caching/fetch process it uses.
  • Each chunk process a part of the original data taking into account a padding margin to avoid border effects. From my experience this is enough and overlapping is not necessary. Also that was the thing forcing us to create the division mask.
  • Each chunk is then processed by the model.
  • The prediction is processed by the specific workflow at hand via after_one_patch_prediction_by_chunks function. In the old implementation, this process was done in a subsequent step after the whole image was predicted. Now, with the generator this part should be faster I think.

Notes:
Now, a Zarr per each rank (GPU) is created first and then the main rank gathers all into just one. I tried several times to avoid that gathering step by creating just the final Zarr image instead, i.e. not creating one per each rank, but there is ALWAYS a data loss. From the old zarr version docs you can read that if the data is chunked with no overlapping chunks, so do I, it should be possible to write simultaneously in the Zarr. They provide also some sync through file/threads. All options lead to data loss.
I tried using latest zarr version (3.0.6) but it is already not documented how to do this parallel writing plus I faced several errors when adapting the code to the new zarr usage. I think they are still solving issues from the migration from Zarr 2 to Zarr 3: maybe we need to wait until they document that so we can try it again.

On the other hand, now you can filter out those regions of the test data based on some statistics. I think you should try it out to reduce the number of patches that are processed, e.g. to discard all background patches. Find the variables here:

BiaPy/biapy/config/config.py

Line 538 in 483c98e

_C.DATA.TEST.FILTER_SAMPLES = CN()

. I think it is worth trying all these things with a medium size file, so you can adjust the filtering, and then apply the model to the whole huge image(s).

Let me know what do you think!

[ClementCaporal]

Hi @danifranco, (You cooked exactly 2 weeks ahah)

Fantastic, you solve many problems with thoses changes. I will try when I have time to check the time effect on our dataset.
I will also check the border effect with padding without overlapping.

I will check the code but in case it is simplier for you to answer directly:

• when the patch is point detection is done subsequently, is the padding from the patch already removed?
• Is the detection conducted by another thread? If so, can we choose the number of worker for the inference / detection ?
• Filter sample will be super useful. Does it run on a specific zarr layer ? (ie, can I choose to check on layer 2 to skip a batch of layer 0 patches?)

Thank you a lot for the update!

[danifranco]

Hi,

Ofc, I'm a man of my word! ;)

when the patch is point detection is done subsequently, is the padding from the patch already removed?

Yeah, the padding is already removed by that time. See when after_one_patch_prediction_by_chunks is called and how above the padding is removed here:

BiaPy/biapy/engine/base_workflow.py

Line 1381 in e19664f

self.after_one_patch_prediction_by_chunks(single_pred, patch_in_data[i])

Is the detection conducted by another thread? If so, can we choose the number of worker for the inference / detection ?

The loader gave a patch to each GPU/rank. Each rank then predicts and after that processes the detection on that patch (without padding). So this process is done by each rank, which means is it multithreaded always. There are no two steps now, everything is done together.

Filter sample will be super useful. Does it run on a specific zarr layer ? (ie, can I choose to check on layer 2 to skip a batch of layer 0 patches?)

I'm not understanding this very well. You mean if it can work on specific channels of the data? it can not, each patch is processed entirely to determine if it is discarded or not.

[ClementCaporal]

Yeah, the padding is already removed by that time. See when after_one_patch_prediction_by_chunks is called and how above the padding is removed here:

I would argue that the padding should be kept so the detection can run smoothly. Because I think what we discussed here is still valid: #50

I'm not understanding these very well. You mean if it can work on specific channels of the data? it can not, each patch is processed entirely to determine if it is discarded or not.

In my experience, loading patch can be the time bottleneck. So the less patch we load the better. When I am writing layer I thought of "zarr level resolution". Loading a patch from a lower resolution zarr ( my_zarr/1 vs my_zarr/2 in ome.zarr) would allow me to skip faster patches --> faster inference!

---

a Zarr per each rank (GPU) is created first and then the main rank gathers all into just one.

I could also add a PR that do not right the inference at all?
The user (at least for me) workflow could be:

1. check on a subpart of the To what the inference looks like.
2. on the whole dataset without saving the inference

[danifranco]

I would argue that the padding should be kept so the detection can run smoothly. Because I think what we discussed here is still valid: #50

I don't think this problem still applies. Try with the implementation that we have now and if there is still such problem then we could leave the padding before process.

In my experience, loading patch can be the time bottleneck. So the less patch we load the better.

Greate cause the filters are applied before feding the patches to the GPUs so it will be an speed up here too.

When I am writing layer I thought of "zarr level resolution". Loading a patch from a lower resolution zarr ( my_zarr/1 vs my_zarr/2 in ome.zarr) would allow me to skip faster patches --> faster inference!

Oh, I see. You can decide were the data is stored within the Zarr enabling DATA.TEST.INPUT_ZARR_MULTIPLE_DATA and using DATA.TEST.INPUT_ZARR_MULTIPLE_DATA_RAW_PATH por the path. Find these variables in the config.py to see more details.

I could also add a PR that do not right the inference at all? The user (at least for me) workflow could be:

1. check on a subpart of the To what the inference looks like.
2. on the whole dataset without saving the inference

I'm not getting this. If you want to just check how the inference is done you can try with a small zarr image.

[ClementCaporal]

Thank you for your answers!

I could also add a PR that do not right the inference at all? The user (at least for me) workflow could be:

    check on a subpart of the To what the inference looks like.
    on the whole dataset without saving the inference

I'm not getting this. If you want to just check how the inference is done you can try with a small zarr image.

If I understand correctly in the new detection workflow
load -> inference per patch + padding -> detect per patch -> save inference in .zarr (per rank) and detection .csv (per patch)

I think it can be useful to add an option to remove saving the .zarr to save only the .csv

[danifranco]

load -> inference per patch + padding -> detect per patch -> save inference in .zarr (per rank) and detection .csv (per patch)

There is no .csv per patch now, only the final one is generated.

Now I'm understanding you. I get your point of avoid creating the zarr parts and final zarr. Feel free to add a PR or give me a few days to do it.