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

Question

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

感谢您的帮助。

Answer 1

我根据 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

更新 2019.10.06：增加了对包含其他模型作为层的模型的支持。

UPDATE 2020.07.17：函数现在可以在 TensorFlow v2 中正常工作。

Answer 2

希望对你有帮助...

• 这里是如何确定 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)]))

Answer 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

Answer 4

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