Argmax in Sliding Window Inference process_fn

[landynfrancis]

Hello all,

I am trying to figure out how to perform sliding_window_inference on large 3D images without running out of GPU memory. I have a collection of 4090's and have already figured out how to use torch.nn.DataParallel to perform the actual inference on multiple GPU's. My issue comes with performing the argmax on the inferred image. The 4090 can't handle the additional storage needed to compute the argmax, and doing it on the CPU is very slow.

So that brings us to my idea:

I am trying to use the process_fn argument in sliding_window_inference to perform the argmax on the individual windows before they are patched back together.

I could use some advice on how to accomplish this, also if it is a bad idea, any ideas to solve my argmax memory issue would be appreciated.

Thanks