如何确定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ícioPereira的回答创建了完整的函数.

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

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

UPD 2019.10.06 :添加了对包含其他模型作为图层的模型的支持.

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

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