Keras (2.0.8) с TensorFlow (1.3) backend занимает всю доступную оперативную память

Я использую библиотеку keras с включенным бэкендом tensorflow и CUDA. Пакеты типун выхода версии:

Keras (2.0.8)

tensorflow-gpu (1.3.0)

tensorflow-tensorboard (0.1.8)

У меня есть следующий код, который создает VGG16 модель и загружает веса ImageNet:

def create_vgg16_model(target_size: tuple, n_classes: int):

    base = VGG16(include_top=False,

                 input_shape=target_size,

                 weights='imagenet')



    x = base.output

    x = Flatten()(x)

    x = Dense(n_classes, activation='softmax', name='top')(x)



    model = Model(inputs=base.input, outputs=x)

    for layer in model.layers[:-1]:

        layer.trainable = False



    model.compile(optimizer='adam', loss='categorical_crossentropy')

    return model

Обучение модели идет хорошо, и nvidia-smi показывает, что память GPU используется по мере необходимости. Но затем я проверил вывод команды top и вот что я вижу:

  PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND                                                                 

 1268 ck        20   0  166288  31964  12416 S  29.5  0.1  13:05.39 Xtightvnc                                                               

32235 ck        30  10   32252   3700   3356 S   5.3  0.0   0:36.93 cwaves 

------------------------------------------------------------------------------    

32212 ck        20   0 27.485g 1.184g 190404 S   2.3  3.8   0:35.44 python  

------------------------------------------------------------------------------                                                                

26015 root      20   0       0      0      0 S   0.3  0.0   0:00.30 kworker/3:1                                                             

31754 ck        20   0   43168   3904   3080 R   0.3  0.0   0:04.45 top                                                                     

    1 root      20   0  185644   6204   3984 S   0.0  0.0   0:10.44 systemd                                                                 

Я прошелся по коду с отладчиком и понял, что память выделена в следующая функция, взятая из keras.backend.tensorflow_backed, которая создает объект tf.Session:

def get_session():        

    global _SESSION

    if tf.get_default_session() is not None:

        session = tf.get_default_session()

    else:

        if _SESSION is None:

            if not os.environ.get('OMP_NUM_THREADS'):

                config = tf.ConfigProto(allow_soft_placement=True)

            else:

                num_thread = int(os.environ.get('OMP_NUM_THREADS'))

                config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,

                                        allow_soft_placement=True)

                # next line allocates ~28GB of RAM

                _SESSION = tf.Session(config=config)

        session = _SESSION

    if not _MANUAL_VAR_INIT:

        with session.graph.as_default():

            _initialize_variables()

    return session

И это происходит для всех доступных моделей, потому что память выделяется при создании сеанса, до начала обучения или инициализации переменных.

Я знаю, что TF выделяет всю доступную память GPU (если вы не переопределяете ConfigProto и/или не настраиваете переменные среды), но делает ли он то же самое с ОЗУ? То есть, похоже, что фреймворк выделяет всю оперативную память, которая у меня есть на моем компьютере. машина, кроме той, которая уже выделена другими процессами.

Кто-нибудь заметил такое поведение с различными версиями tensorflow или keras? Как вы думаете, есть ли способ каким-то образом ограничить объем используемой памяти?

---

Обновление 1

Некоторое время назад один из моих обучающих скриптов был убит ядром с ошибкой out-of-memory после 50-60 тренировочных эпох. Хотя статистика использования энергозависимой памяти GPU показывает, что она также используется. (Не только выделено, как у меня понятно).

---

Обновление 2

Согласитесь, виртуальная память не является допустимой метрикой. Но я выяснил, что потребление памяти почти линейно увеличивается в процессе обучения модели. У меня есть следующий цикл обучения:

def train_model(model, x, y):

    loss = model.train_on_batch(x, y)

    return loss





def train_model_42(model, x, y):

    # dummy function

    return 42.0





def training_loop():

    # training parameters

    target_size = 224, 224, 3

    batch_size = 128



    # generator yielding batches of file paths

    files_stream = FilesStream(folder=TRAIN_IMAGES, batch_size=batch_size)

    files_source = files_stream()



    # list of generators loading images from persistent storage

    gens = [

        image_loader(),

        augment_images(horizontal_flip=True),

        shuffle_samples(),

        normalize_images(target_size=target_size)

    ]



    # Model from keras.applications with replaced top layer

    model = get_model('resnet50').build(n_classes=n_classes)



    for epoch in range(1, 1001):

        epoch_loss = []

        for _ in range(files_stream.steps_per_epoch):

            for gen in gens:

                gen.send(None)

            processed = next(files_source)

            for gen in gens:

                processed = gen.send(processed)

            x, y = processed

            loss = train_model_42(model, x, y) # <-- this shows pic. 1

            # loss = train_model(model, x, y)    <-- this shows pic. 2                  

            epoch_loss.append(loss)

        avg_loss = sum(epoch_loss) / len(epoch_loss)

        print('Epoch %03d: train loss = %2.4f' % (epoch, avg_loss))

Когда я использую фиктивную обучающую функцию, график потребления памяти выглядит так, как показано на pic 1:

Но при выполнении реального тренировочного процесса, это выглядит как pic 2:

Почему потребление памяти увеличивается в процессе обучения? Кэшируются ли предыдущие пакеты данных или что-то еще? Должна ли модель / веса или что-то еще занимать все больше и больше памяти?

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