How to calculate the mini-batch memory impact when training deep learning models?

Question

I'm trying to calculate the amount of memory needed by a GPU to train my model based on this notes from Andrej Karphaty.

My network has 532,752 activations and 19,072,984 parameters (weights and biases). These are all 32 bit floats values, so each takes 4 bytes in memory.

My input image is 180x50x1 (width x height x depth) = 9,000 float 32 values. I don't use image augmentation, so I think the miscellaneous memory would be only related to the mini-batch size. I'm using a mini-batch size of 128 images.

Based on Andrej's recommendation, I get the following memory sizes:

Activations: 532,752*4/(1024^2) = 2.03 MB

Parameters: 19,072,984*4/(1024^2) * 3 = 218.27 MB

Miscellaneous: 1289,0004/(1024^2) = 4.39 MB

So the total memory to train this network would be 224,69 MB.

I'm using TensorFlow and I think I'm missing something. I haven't run the training yet, but I'm pretty sure (based on past experiences) that the memory in use will be much higher than what I've calculated.

If for each image in the mini-batch, TensorFlow keeps their gradients so it can normalize them later for a single weights/biases updates step, then I think the memory should take into account another 532,752 * 128 values (gradients for each image in the mini-batch). If that is the case, then I'd need more 260.13 MB to train this model with 128 images/mini-batch.

Can you help me understand the memory considerations for training my deep learning model? Are the above considerations right?

Answer 1

@StatsSorceress TL;DR:

I'm going through this activity to see if I can calculate the memory required myself:

Activations: 532,752 * 2 * 4 / (1024^2) = 4.06 MB

Parameters: 19,072,984 * 4 / (1024^2) * 3 = 218.27 MB

Miscellaneous: 128 * 9,000 * 4 / (1024^2) = 4.39 MB

Total Memory: (4.06 * 128) + 218.27 + 4.39 = 742.34 MB

(Someone please correct me on this if I'm wrong. FYI, you already multiplied miscellaneous by 128, so that's why I didn't multiply it by 128 above)

---

I would point you to this article and the corresponding video. They helped me to understand what is going on a lot better.

NOTE: The memory required to use a network for predictions is far less than that required for training for two reasons:

• When predicting, we only send an image forward through the network and not backward (so we don't multiply memory X 3; see below)
• There's one prediction per image (so we don't need to multiply the memory required for one image by a batch size because we don't use batches in prediction).

Process (Memory to Train)

1. Calculate the memory required to train on one image
2. Multiply this number by the number of images in your batch

(REMEMBER: Mini-batching says we take a subset of our data, compute the gradients and errors for each image in the subset, then average these and step forward in the direction of the average. For convnets, weights and biases are shared, but the number of activations is mutliplied by the number of images in the batch.).

STEP 1: Memory for 1 Image

To train one image, you must reserve memory for:

• Model Parameters:

  The weights and biases at each layer, their gradients, and their momentum variables (if Adam, Adagrad, RMSProp, etc., optimizers are used)

  To approximate the memory for this, calculate the memory required to store the weights and biases and multiply that by 3 (i.e. "by 3" because we're saying the amount of memory needed to store the weights and biases is (roughly) equal to that needed for the gradients and for the momentum variables)

  EQUATIONS:

  Convolutions:

  weights(n) = depth(n) * (kernel_width * kernel_height) * depth(n-1)

  biases(n) = depth(n)

  Fully Connected (Dense) Layers:

  weights(n) = outputs(n) * inputs(n)

  biases(n) = outputs(n)

where n is the current layer and n-1 is the previous layer, and outputs are the number of outputs from the FC layer and inputs are the number of inputs to the FC layer (if the previous layer is not a fully-connected layer, the number of inputs is equal to the size of that layer flattened).

NOTE: The memory for the weights and biases alone, plus the memory for the activations for one image (see below), is the total amount of memory you need for predictions (excluding some overhead for memory for convolutions and some other things).

• Activations (these are "Blobs" in Caffe):

(I'm using terms loosely here, bear with me)

Each convolution in a convolution layer produces "number of pixels in image" activations (i.e. you pass an image through a single convolution, you get a single feature map consisting of "m" activations, where "m" is the number of pixels from your image/input).

For fully-connected layers, the number of activations you produce is equal to the size of your output.

Convolutions:

activations(n) = image_width * image_height * image_num_channels

Fully Connected (Dense) Layers:

activations(n) = outputs(n)

Note that your input is really only an image at the beginning of the network. After convolutions, it turns into something else (feature maps). So really replace "image_width", "image_height", and "image_num_channels" with "input_width", "input_height", and "layer_depth" to be more precise. (It's just easier for me to think of this concept in terms of images.)

Since we also need to store the error for the activations at each layer (used in the backward pass), we multiply the number of activations by 2 to get the total number of entities we need to make room for in our storage space. The number of activations increases with the number of images in the batch, so you multiply this number by the batch size.

STEP 2: Memory to Train Batch

Sum the number of weights and biases (times 3) and the number of activations (times 2 times the batch size). Multiply this by 4, and you get the number of bytes required to train the batch. You can divide by 1024^2 to get the answer in GB.

Answer 2

I think you're on the right track.

Yes, you will need to store the derivatives of the activations and of the parameters for backpropagation.

Additionally, your choice of optimization may matter. Are you training using SGD, or Adam, or Adagrad? These will all have different memory requirements. For example, you're going to have to store the step size cache for a momentum-based method, although that should be secondary compared to the other memory considerations you mention.

So all in all, you seem to have calculated the memory requirements for a forward pass. Andrej Karpathy mentions that the backward pass could take up to 3x the memory of the forward pass, so this might be why you see such a difference (scroll down to 'Case Studies' on the website to see an example for VGGNet).

Answer 3

Alternatively, I think you can use any profiler library to analyze the memory and CPU usage by your program. There are many python libraries which can give you snapshot of memory and CPU usage by particular thread or process at millisecond interval.

You may run the part of your program you want to monitor in a different sub-process using popen and monitor it's memory and CPU usage using it's PID.

psutil I find good for such work. Though there are many others.

I hope this will help.