Оценка максимального размера партии в Керасе

Question

Я пишу функцию для оценки максимального размера пакета, который может вписаться в графический процессор.Чтобы проверить мое решение, я написал функцию, которая генерирует много случайных вычислительных графиков из основных слоев Keras и компилирует модели Keras.Затем я передаю модели в эту функцию:

from itertools import chain
from math import log, floor

import keras.backend as K
import operator as op
from functools import reduce
from keras.models import Model


def estimate_batch_size(model: Model, available_mem: int,
                        scale_by: float = 5.0,
                        precision: int = 2) -> int:
    """
    :param model: keras Model
    :param available_mem: available memory in bytes
    :param scale_by: scaling factor
    :param precision: float precision: 2 bytes for fp16, 4 - for fp32, etc.
    :return: closest 2^n to the estimated batch size
    """

    num_params = sum(chain.from_iterable((
        (reduce(op.mul, l.output_shape[1:]) for l in model.layers),
        (K.count_params(x) for x in model.trainable_weights),
        (K.count_params(x) for x in model.non_trainable_weights)
    )))
    max_size = int(available_mem / (precision * num_params * scale_by))
    return int(2 ** floor(log(max_size, 2)))

Я добавил коэффициент масштабирования, чтобы получить консервативную оценку и избежать ошибок «нехватки памяти» (OOM) во время автоматической оптимизации гиперпараметров.Так или иначе, даже с коэффициентом масштабирования 4,5, я получаю ООМ.Вот примерная сводка модели

______________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
input_252 (InputLayer)       (None, 50, 45, 1)         0
_________________________________________________________________
dense_696 (Dense)            (None, 50, 45, 34)        68
_________________________________________________________________
dense_699 (Dense)            (None, 50, 45, 279)       9765
=================================================================
Total params: 9,833
Trainable params: 9,833
Non-trainable params: 0
_________________________________________________________________
None

Здесь с учетом 8 ГБ видеопамяти и scale_by=4.5 функция возвращает размер пакета 1024 (параметры 716333 fp32, включая промежуточные заполнители, умноженные на 4,5).Тем не менее, несмотря на большой коэффициент масштабирования, я все еще получаю ООМ.Я знаю, что этот метод, среди прочего, не учитывает заполнители, выделенные для вычислений градиента, но я все еще удивлен тем фактом, что даже коэффициент масштабирования 4,5 не дает надежной оценки.Можно ли получить более точную оценку?