0
$\begingroup$

I was trying to estimate the memory usage for my LSTM network by referring to an examples of CNN memory usage calculation at http://cs231n.github.io/convolutional-networks/#computational-considerations . The LSTM network architecture is as follows:enter image description here

My input data is of dimension (None,4,34569) float 32 values and With batchsize=8, the following are my estimations of memory usage

Parameters: 482646529*(4/2^20)*3=5.4GB

Given as No_of_parameters *memory_per_float32*factor.

The factor of 3 is to consider the memory to hold parameters, gradients in BP and optimization cache(if adam, RMSprop etc are used).

Activations: No_of_activations*(4/2^20)x2 =129857*(4/2^20)*2 =1MB

The factor of 2 since the memory must hold the activations for backward pass

Here the No_of_activations are calculated as follows:

lstm_1: 6*4*3072=73728

(used '6' because of forget,input,output gate, cell state, cell candidate,output activation and '4' for the length of time sequence)

dropout_1: 4*3072=12288

lstm_2: 6*4*1024=24576

dropout_2: 4*1024=4096

lstm_3: 6*4*512=12288

dropout_3: 4*512=2048

time_distributed_Dense_1: 256

dropout_4 :256

time_distributed_Dense_2 : 128

dropout_5: 128

time_distributed_Dense_3: 64

time_distributed_Dense_4: 1

Total number of activations: 129857

Miscellaneous : 4*34569*(4/2^20)=0.52MB

Total Memory :parameters_memory+activations * batchsize+Miscellaneous * batchsize = 5.4GB + 1MB *8 +0.52MB *8 ~=5.4GB

I am not sure if the no of activations required has been calculated correctly. Also Please let me know if the above memory calculations are appropriate. I am trying to train this network on 12GB GPU but I am running out of resource in the first epoch. If the above estimates are not appropriate, then could you please hint me towards a way to find roughly the memory usage for an LSTM network?

$\endgroup$
  • $\begingroup$ did you end up figuring this out? $\endgroup$ – Daniel V Apr 7 at 23:56

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.