如何确定 Keras 模型所需的内存? [英] How to determine needed memory of Keras model?

问题描述

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

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?

感谢您的帮助.

推荐答案

我根据 Fabrício Pereira 的回答创建了一个完整的函数.

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:添加了对包含其他模型作为层的模型的支持.

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

UPDATE 2020.07.17:函数现在可以在 TensorFlow v2 中正常工作.

UPDATE 2020.07.17: Function now works correctly in TensorFlow v2.

这篇关于如何确定 Keras 模型所需的内存?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！