How to calculate the GPU memory need to run a deep learning network?

Question

In general, how do I calculate the GPU memory need to run a deep learning network?

I'm asking this question because my training for some network configuration is getting out of memory.

If the TensorFlow only store the memory necessary to the tunable parameters, and if I have around 8 million, I supposed the RAM required will be:

RAM = 8.000.000 * (8 (float64)) / 1.000.000 (scaling to MB)

RAM = 64 MB, right?

The TensorFlow requires more memory to store the image at each layer?

By the way, these are my GPU Specifications:

• Nvidia GeForce 1050 4GB

Networking topology

• Unet
• Input Shape (256,256,4)

Model: "functional_1"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 256, 256, 4) 0                                            
__________________________________________________________________________________________________
conv2d (Conv2D)                 (None, 256, 256, 64) 2368        input_1[0][0]                    
__________________________________________________________________________________________________
dropout (Dropout)               (None, 256, 256, 64) 0           conv2d[0][0]                     
__________________________________________________________________________________________________
conv2d_1 (Conv2D)               (None, 256, 256, 64) 36928       dropout[0][0]                    
__________________________________________________________________________________________________
max_pooling2d (MaxPooling2D)    (None, 128, 128, 64) 0           conv2d_1[0][0]                   
__________________________________________________________________________________________________
conv2d_2 (Conv2D)               (None, 128, 128, 128 73856       max_pooling2d[0][0]              
__________________________________________________________________________________________________
dropout_1 (Dropout)             (None, 128, 128, 128 0           conv2d_2[0][0]                   
__________________________________________________________________________________________________
conv2d_3 (Conv2D)               (None, 128, 128, 128 147584      dropout_1[0][0]                  
__________________________________________________________________________________________________
max_pooling2d_1 (MaxPooling2D)  (None, 64, 64, 128)  0           conv2d_3[0][0]                   
__________________________________________________________________________________________________
conv2d_4 (Conv2D)               (None, 64, 64, 256)  295168      max_pooling2d_1[0][0]            
__________________________________________________________________________________________________
dropout_2 (Dropout)             (None, 64, 64, 256)  0           conv2d_4[0][0]                   
__________________________________________________________________________________________________
conv2d_5 (Conv2D)               (None, 64, 64, 256)  590080      dropout_2[0][0]                  
__________________________________________________________________________________________________
max_pooling2d_2 (MaxPooling2D)  (None, 32, 32, 256)  0           conv2d_5[0][0]                   
__________________________________________________________________________________________________
conv2d_6 (Conv2D)               (None, 32, 32, 512)  1180160     max_pooling2d_2[0][0]            
__________________________________________________________________________________________________
dropout_3 (Dropout)             (None, 32, 32, 512)  0           conv2d_6[0][0]                   
__________________________________________________________________________________________________
conv2d_7 (Conv2D)               (None, 32, 32, 512)  2359808     dropout_3[0][0]                  
__________________________________________________________________________________________________
conv2d_transpose (Conv2DTranspo (None, 64, 64, 256)  524544      conv2d_7[0][0]                   
__________________________________________________________________________________________________
concatenate (Concatenate)       (None, 64, 64, 512)  0           conv2d_transpose[0][0]           
                                                                 conv2d_5[0][0]                   
__________________________________________________________________________________________________
conv2d_8 (Conv2D)               (None, 64, 64, 256)  1179904     concatenate[0][0]                
__________________________________________________________________________________________________
dropout_4 (Dropout)             (None, 64, 64, 256)  0           conv2d_8[0][0]                   
__________________________________________________________________________________________________
conv2d_9 (Conv2D)               (None, 64, 64, 256)  590080      dropout_4[0][0]                  
__________________________________________________________________________________________________
conv2d_transpose_1 (Conv2DTrans (None, 128, 128, 128 131200      conv2d_9[0][0]                   
__________________________________________________________________________________________________
concatenate_1 (Concatenate)     (None, 128, 128, 256 0           conv2d_transpose_1[0][0]         
                                                                 conv2d_3[0][0]                   
__________________________________________________________________________________________________
conv2d_10 (Conv2D)              (None, 128, 128, 128 295040      concatenate_1[0][0]              
__________________________________________________________________________________________________
dropout_5 (Dropout)             (None, 128, 128, 128 0           conv2d_10[0][0]                  
__________________________________________________________________________________________________
conv2d_11 (Conv2D)              (None, 128, 128, 128 147584      dropout_5[0][0]                  
__________________________________________________________________________________________________
conv2d_transpose_2 (Conv2DTrans (None, 256, 256, 64) 32832       conv2d_11[0][0]                  
__________________________________________________________________________________________________
concatenate_2 (Concatenate)     (None, 256, 256, 128 0           conv2d_transpose_2[0][0]         
                                                                 conv2d_1[0][0]                   
__________________________________________________________________________________________________
conv2d_12 (Conv2D)              (None, 256, 256, 64) 73792       concatenate_2[0][0]              
__________________________________________________________________________________________________
dropout_6 (Dropout)             (None, 256, 256, 64) 0           conv2d_12[0][0]                  
__________________________________________________________________________________________________
conv2d_13 (Conv2D)              (None, 256, 256, 64) 36928       dropout_6[0][0]                  
__________________________________________________________________________________________________
conv2d_14 (Conv2D)              (None, 256, 256, 1)  65          conv2d_13[0][0]                  
==================================================================================================
Total params: 7,697,921
Trainable params: 7,697,921
Non-trainable params: 0

This is the error given.

---------------------------------------------------------------------------
ResourceExhaustedError                    Traceback (most recent call last)
<ipython-input-17-d4852b86b8c1> in <module>
     23 # Train the model, doing validation at the end of each epoch.
     24 epochs = 30
---> 25 result_model = model.fit(train_gen, epochs=epochs, validation_data=val_gen, callbacks=callbacks)

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\keras\engine\training.py in _method_wrapper(self, *args, **kwargs)
    106   def _method_wrapper(self, *args, **kwargs):
    107     if not self._in_multi_worker_mode():  # pylint: disable=protected-access
--> 108       return method(self, *args, **kwargs)
    109 
    110     # Running inside `run_distribute_coordinator` already.

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\keras\engine\training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
   1096                 batch_size=batch_size):
   1097               callbacks.on_train_batch_begin(step)
-> 1098               tmp_logs = train_function(iterator)
   1099               if data_handler.should_sync:
   1100                 context.async_wait()

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\def_function.py in __call__(self, *args, **kwds)
    778       else:
    779         compiler = "nonXla"
--> 780         result = self._call(*args, **kwds)
    781 
    782       new_tracing_count = self._get_tracing_count()

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\def_function.py in _call(self, *args, **kwds)
    838         # Lifting succeeded, so variables are initialized and we can run the
    839         # stateless function.
--> 840         return self._stateless_fn(*args, **kwds)
    841     else:
    842       canon_args, canon_kwds = \

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\function.py in __call__(self, *args, **kwargs)
   2827     with self._lock:
   2828       graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
-> 2829     return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
   2830 
   2831   @property

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\function.py in _filtered_call(self, args, kwargs, cancellation_manager)
   1846                            resource_variable_ops.BaseResourceVariable))],
   1847         captured_inputs=self.captured_inputs,
-> 1848         cancellation_manager=cancellation_manager)
   1849 
   1850   def _call_flat(self, args, captured_inputs, cancellation_manager=None):

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1922       # No tape is watching; skip to running the function.
   1923       return self._build_call_outputs(self._inference_function.call(
-> 1924           ctx, args, cancellation_manager=cancellation_manager))
   1925     forward_backward = self._select_forward_and_backward_functions(
   1926         args,

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\function.py in call(self, ctx, args, cancellation_manager)
    548               inputs=args,
    549               attrs=attrs,
--> 550               ctx=ctx)
    551         else:
    552           outputs = execute.execute_with_cancellation(

~\Anaconda3\envs\tf23\lib\site-packages\tensorflow\python\eager\execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58     ctx.ensure_initialized()
     59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:
     62     if name is not None:

ResourceExhaustedError:  OOM when allocating tensor with shape[8,64,256,256] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[node gradient_tape/functional_1/conv2d_14/Conv2D/Conv2DBackpropInput (defined at <ipython-input-17-d4852b86b8c1>:25) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
 [Op:__inference_train_function_17207]

Function call stack:
train_function

Is there any type of mistake in the network definition? How could I improve the network to solve this problem?

Answer 1

In fact, I do not know how to calculate GPU memory to run a neural network but I have a solution for allocation problems in GPUs while using tensorflow framework.

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    # Restrict TensorFlow to only allocate 2GB * 2 of memory on the first GPU
    try:
        tf.config.experimental.set_virtual_device_configuration(
            gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048 * 2)])
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Virtual devices must be set before GPUs have been initialized
        print(e)

You can set a memory limit on GPU which sometimes solves memory allocation problems. As shown above, you can set "memory_limit" parameter as your configuration requires.

Also be careful about using correct framework. If you want to use above code to set memory, you have to build your neural network from tensorflow with keras backend.

from tensorflow.python.keras.models import Sequential

Answer 2

You can calculate the memory requirement analytically, but it's still not going to beat physical test in practice as there are so many unknown variables in the system which can takes the GPU memory. Maybe tensorflow will decide to store the gradients, then you have to take into account the memory usage of it also.

The way I do it is by setting the GPU memory limit to a high value e.g. 1GB, then test the model inference speed. Then I repeat the process with half the memory. I do it until the model refuses to run or the model speed drops. For example, I start with 1GB, then 512MB, then 256MB, eventually I got to 32 MB and the model speed drops. At 16MB, the model refuses to run. So I know that 64 MB is the minimum requirement I should use for my model. If I want to get a more precise number, I'd repeat the binary search process a couple more time between 64 MB and 32 MB.

You can see how to limit the GPU memory here: https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)

Answer 3

TensorFlow is interesting that it can store not only weights, but also training data in video RAM.

with tf.device('/gpu:0'):
    tensorflow_dataset = tf.constant(numpy_dataset)

Feeding training data and weights to GPU for matrix mul is faster than from regular RAM.

Video RAM required = Number of params * sizeof(weight type) +
                     Training data amount in bytes

However, I believe that video RAM required should be at least 1.5 times the above value just to be sure things would be working.