Does patch size affect GPU memory consumption?

[lanslotttTT]

Thanks for your great job! I was wondering that the relationship between medical image patch size and GPU consumption during training, is there have any classical papers? Thanks in advance!

[Nic-Ma]

Hi @ericspod , do you know any related papers?

Thanks in advance.

[atbenmurray]

Hi @lanslotttTT. It isn't really a question of papers particularly; neural networks are made out of common building blocks whose memory requirements can be concretely understood; I'll can make a few observations that may help.

One important factor is whether your network is purely convolutional or not. Take a UNet vs. a VAE, for example:

• The UNet is purely convolutional. If you increase the volume of a patch by the ratio that you decrease the batch size, it should require an identical amount of memory
• The VAE has a number of linear layers that take the lowest spatial resolution and "de-cohere" it to a latent space that is then compressed, and then perform the opposite operation for the decoder. The VAE parameter memory requirement is then either linear or super-linear in terms of patch size for certain layers, due to the number of weights such a linear layer requires for back propagation (I think I'm correct but it's late in the evening; I'll circle back to this tomorrow and give a concrete example).

You can examine the effect of changing such parameters on a given network architecture using torchinfo, which allows you to print out the network layers and parameters.

[atbenmurray]

Hi @lanslotttTT, an addendum to make things a little more concrete.

(Note I'm excluding bias parameters from the calculation because they are trivially small)
Assuming 2D spatial data:

• A Conv2D layer requires [c_in, c_out, k_size, k_size] parameters. This is independent of image size
• A Linear layer requires [c_in, c_out] parameters.
  • If c_in and c_out are independent of image size, parameters are independent of image size
  • If c_in xor c_out are proportional to image size, parameters change proportional to image size
  • If c_in and c_out are proportional to image size, parameters change according to the square of the image size

Note that learnable parameters are just one source of gpu memory usage in a neural network, but you can see that increasing patch_size and decreasing batch size such that the product of these two values is the same can increase the learnable parameters of the network, depending on what is in your network.