GridPatch with pin_memory=True significant slow-down in following epochs

[drbeh]

Reported by a user (@kenza-bouzid) here:

When using num_workers>1 and pin_memory=True, training time increases exponentially over epochs


I had profiled/timed all intermediate steps, and found out that GridPatch was the guilty one

So I timed all intermediate step in the tranform

It turned out that the operation that was taking too long was a memory allocation by np.array

MONAI/monai/transforms/spatial/array.py

Line 3279 in a2ec375

patched_image = np.array(patches[0])

I eventually fixed it by setting pin_memory=False,
which I explain by cuda memory allocation being more expensive as it has to be allocated from the pinned memory

Note that I am dealing with particularly large slides ~50kx50k

Any thoughts on this?

[drbeh]

@kenza-bouzid just to double confirm that when you set pin_memory=False this issue goes away and it takes the same time (few seconds) in all epochs, right?

[kenza-bouzid]

@drbeh Correct! I don't have a proper explanation why pin_memory=False fixes this issue tbh.
Training looks healthy with pin_memory=False, epoch time is constant

[drbeh]

@ericspod @Nic-Ma @wyli @rijobro do you have any insights here on why setting pin_memory=True in dataloader causes this behavior?

It seems that these conversion to and from numpy arrays have some effects here but I thought pinning memory is being done in the last step of dataloader so I don't see why it has anything to do with it.
Also there have been some issues with pin memory in pytorch in the past, like this one from years ago, so can it be a pytorch issue?

[ericspod]

I honestly am not sure what the issue here would be. I recall there being other issues we've encountered in the past with pinned memory so I typically avoid it. I can only speculate it could be the pinned segment's lifetime continues after the np.array conversion because of Pytorch's buffer reuse or some other issue in Pytorch or CUDA itself.