如何确定Keras模型所需的内存?

时间:2017-03-31 09:32:39

标签: memory keras

我正在使用Keras 2.0.0,我想在GPU上训练一个包含大量参数的深度模型。 使用太大的图像,我的内存不足(OOM)。 使用太低的图像,模型的准确性将比可能的更差。 因此,我想找到适合我的GPU的最大可能输入尺寸的图像。 在给定模型和输入数据的情况下,是否有任何计算内存的功能(例如,与model.summary()相当)?

感谢您的帮助。

4 个答案:

答案 0 :(得分:26)

我根据FabrícioPereira的回答创建了完整的功能。

def get_model_memory_usage(batch_size, model):
    import numpy as np
    from keras import backend as K

    shapes_mem_count = 0
    for l in model.layers:
        single_layer_mem = 1
        for s in l.output_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 set(model.trainable_weights)])
    non_trainable_count = np.sum([K.count_params(p) for p in set(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)
    return gbytes

答案 1 :(得分:5)

希望这可以帮助你...

  • 以下是确定Keras模型(var <------starting batch-------> Volume in drive C has no label. Volume Serial Number is 0CE1-E926 Directory of C:\temp 06/26/2017 03:17 PM <DIR> . 06/26/2017 03:17 PM <DIR> .. 08/10/2016 09:07 PM <DIR> Clt-Inst 06/07/2016 04:31 PM 56,406,016 splunkforwarder-6.4.1-debde650d26e-x64-release.msi 06/26/2017 02:44 PM 35,922,892 Windows6.0-KB968930-x64.msu 2 File(s) 92,328,908 bytes 3 Dir(s) 418,918,834,176 bytes free <servername> <------ending batch------> )的多种形状以及每个形状单位在内存中占用4个字节的方式:

      

    model

         

    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]))

  • 以下是确定Keras模型(var memory = shapes_count * 4)的一些参数的方法:

      

    model

         

    from keras import backend as K

         

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

答案 2 :(得分:3)

这是我@ZFTurbo答案的变体。它为嵌套的Keras模型,不同的TensorFlow dtype提供了更好的处理,并消除了对NumPy的依赖。我已经在TensorFlow 2.3.0上进行了编写和测试,但可能无法在早期版本中使用。

def keras_model_memory_usage_in_bytes(model, *, batch_size: int):
    """
    Return the estimated memory usage of a given Keras model in bytes.
    This includes the model weights and layers, but excludes the dataset.

    The model shapes are multipled by the batch size, but the weights are not.

    Args:
        model: A Keras model.
        batch_size: The batch size you intend to run the model with. If you
            have already specified the batch size in the model itself, then
            pass `1` as the argument here.
    Returns:
        An estimate of the Keras model's memory usage in bytes.

    """
    default_dtype = tf.keras.backend.floatx()
    shapes_mem_count = 0
    internal_model_mem_count = 0
    for layer in model.layers:
        if isinstance(layer, tf.keras.Model):
            internal_model_mem_count += keras_model_memory_usage_in_bytes(
                layer, batch_size=batch_size
            )
        single_layer_mem = tf.as_dtype(layer.dtype or default_dtype).size
        out_shape = layer.output_shape
        if isinstance(out_shape, 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 = sum(
        [tf.keras.backend.count_params(p) for p in model.trainable_weights]
    )
    non_trainable_count = sum(
        [tf.keras.backend.count_params(p) for p in model.non_trainable_weights]
    )

    total_memory = (
        batch_size * shapes_mem_count
        + internal_model_mem_count
        + trainable_count
        + non_trainable_count
    )
    return total_memory

答案 3 :(得分:-1)

我相信,如果您使用定制编写的数据生成器或利用keras中的某些现有生成器,它将解决您的问题。 当所有加载的数据对系统造成负担时,通常会发生内存错误,而是使用生成器将数据集分解为段,这样您就不会耗尽内存,并且可以在任何系统上进行训练。