Как определить необходимую память модели Keras?

Я работаю с Keras 2.0.0, и я бы хотел обучить глубокую модель с огромным количеством параметров на графическом процессоре. Используя слишком большие изображения, у меня заканчивается память (OOM). Используя слишком низкие изображения, точность модели будет хуже, чем возможно. Поэтому я хотел бы найти максимально возможный размер входных изображений, подходящих для моего графического процессора. Есть ли какая-либо функциональность, вычисляющая память (например, сравнимую с model.summary()) с учетом модели и входных данных?

Я ценю вашу помощь.

Ответ 1

Я создал полную функцию, основанную на ответе Фабрицио Перейры.

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: добавлена поддержка моделей, которые содержат другие модели в качестве слоев.

Ответ 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)]))