How to determine needed memory of Keras model?

Question

I am working with Keras 2.0.0 and I'd like to train a deep model with a huge amount of parameters on a GPU. Using too large images, I'm running out of memory (OOM). Using too low images, the model's accuracy will be worse than possible. Therefore I'd like to find the biggest possible input size of images that fit to my GPU. Is there any functionality calculating the memory (e.g. comparable to model.summary()) given the model and input data?

I appreciate your help.

Answer 1

I created a complete function based on the answer of Fabrício Pereira.

def get_model_memory_usage(batch_size, model):     import numpy as np     try:         from keras import backend as K     except:         from tensorflow.keras import backend as K      shapes_mem_count = 0     internal_model_mem_count = 0     for l in model.layers:         layer_type = l.__class__.__name__         if layer_type == 'Model':             internal_model_mem_count += get_model_memory_usage(batch_size, l)         single_layer_mem = 1         out_shape = l.output_shape         if type(out_shape) is list:             out_shape = out_shape[0]         for s in out_shape:             if s is None:                 continue             single_layer_mem *= s         shapes_mem_count += single_layer_mem      trainable_count = np.sum([K.count_params(p) for p in model.trainable_weights])     non_trainable_count = np.sum([K.count_params(p) for p in model.non_trainable_weights])      number_size = 4.0     if K.floatx() == 'float16':         number_size = 2.0     if K.floatx() == 'float64':         number_size = 8.0      total_memory = number_size * (batch_size * shapes_mem_count + trainable_count + non_trainable_count)     gbytes = np.round(total_memory / (1024.0 ** 3), 3) + internal_model_mem_count     return gbytes 

UPDATE 2019.10.06: Added support for models which contain other models as layers.

UPDATE 2020.07.17: Function now works correctly in TensorFlow v2.

Answer 2

Hope this can help you...

• Here is how determinate a number of shapes of you Keras model (var model), and each shape unit occupies 4 bytes in memory:

  shapes_count = int(numpy.sum([numpy.prod(numpy.array([s if isinstance(s, int) else 1 for s in l.output_shape])) for l in model.layers]))

  memory = shapes_count * 4

• And here is how determinate a number of params of your Keras model (var model):

  from keras import backend as K

  trainable_count = int(numpy.sum([K.count_params(p) for p in set(model.trainable_weights)]))

  non_trainable_count = int(numpy.sum([K.count_params(p) for p in set(model.non_trainable_weights)]))