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

Question

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

我感谢您的帮助.

Answer 1

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

def get_model_memory_usage(batch_size, model):
    import numpy as np
    from 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
        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) + internal_model_mem_count
    return gbytes

Answer 2

这是我对@ZFTurbo 的回答的变体。它为嵌套的 Keras 模型、不同的 TensorFlow dtypes 提供了更好的处理，并消除了对 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

Answer 3

希望这可以帮到你...

• 这是确定Keras模型（var model）的许多形状的方式，每个形状单位在内存中占用4个字节：

  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

• 以下是确定Keras模型（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)]))