MCCachingModule and `replicate_in_memory=False`

[arthur-thuy]

Is your feature request related to a problem? Please describe. Not really a feature request but more of a discussion about the choice of replicate_in_memory in the ModelWrapper when using the new MCCachingModule.

I noticed that the MCCachingModule documentation uses replicate_in_memory=False in the experiments. Before, I always used True in my code, which is also used in experiments/vgg_mcdropout_cifar10.py as True is the default. As such, I compared runtimes on an NVIDIA RTX A5000.

Number of iterations	20	50	100
Regular - False	0:16	0:39	1:18
Regular - True	0:11	0:28	0:55
MCCaching - False	0:09	0:20	0:40
MCCaching - True	0:12	0:30	0:57

Findings:

• For False, MCCaching is around 45%-50% faster than Regular. (This is lower than the 70% found for the GeForce 1060Ti)
• "MCCaching - False" is around 20%-30% faster than "Regular - True"
• For True, MCCaching is around 5% slower than Regular

Do you have an idea why MCCaching and replicate_in_memory have such interaction effects?

Describe the solution you'd like In the documentation, it might be useful to discourage the combination of MCCaching and replicate_in_memory=True as this is 5% slower than the "Regular - True" scenario.

Describe alternatives you've considered /

Additional context /

[Dref360]

Hello, we should definitely explain this better, but the ideal case will always be MCCaching + replicate_in_memory=False.

MCCaching works by caching the output of each layer so that it doesn't recompute during MCDropout.

When we use replicate_in_memory=False, the MC Sampling loop is replaced with:

predictions = [model(input) for _ in range(Iterations)]

This means that we compute a lot of redundant features. When adding MCCaching, we compute the high-level features (before Dropout layers) only once.

I will add a warning in the code directly I think. If there is a caching layer somewhere in the model we should warn that it is not optimized. What do you think?

[arthur-thuy]

Hello, thank you for the response. I agree that a warning in the code is most useful. In the documentation, a brief explanation of how the module works internally might indeed also encourage users to try it out.