Tensorflow (Keras) и многопроцессорность приводят к нехватке памяти GPU - PullRequest
Новая
       11

Tensorflow (Keras) и многопроцессорность приводят к нехватке памяти GPU

1 голос
/ 21 марта 2019

У меня есть собственный DataGenerator, который использует модуль многопроцессорной обработки Python для генерации обучающих данных, которые подаются в модель Tensorflow.

Проблема заключается в том, что всякий раз, когда инициализируется новый процесс DataGenerator, создается впечатление, что он пытается инициализировать Tensorflow (который импортируется поверх кода) и выделяет для себя некоторую память графического процессора.

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

Новые процессы и код Tensorflow инициируются в том же файле Python. Есть ли правильный способ использовать многопроцессорность, не позволяя порожденным процессам импортировать Tensorflow и выделять часть памяти GPU для себя?

Вот часть кода (запускается в Windows) для пояснения: 

from multiprocessing import Process, Queue
from multiprocessing.pool import Pool

import cv2
import numpy as np
import tensorflow as tf

from keras.models import load_model

def TrainQueueProcess(queue):
    # This Function Fills The Queue For Other Consumers

def get_model(model_path=None):
    import tensorflow as tf
    import keras.backend.tensorflow_backend as ktf

    def get_session(gpu_fraction=0.333):
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction, allow_growth=True)
        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

    ktf.set_session(get_session())

    from keras import Input, Model
    from keras.applications.mobilenetv2 import MobileNetV2
    from keras.layers import Dense, Dropout
    from keras.optimizers import adam
    from keras.utils import plot_model

    input_tensor = Input(shape=(128, 128, 3))
    base_model = MobileNetV2(weights='imagenet', include_top=False, input_tensor=input_tensor, input_shape=(128, 128, 3), pooling='avg')
    for layer in base_model.layers:
        layer.trainable = True

    op = Dense(128, activation='relu')(base_model.output)
    op = Dropout(.25)(op)
    output_tensor = Dense(2, activation='softmax')(op)
    model = Model(inputs=input_tensor, outputs=output_tensor)
    model.compile(optimizer=adam(lr=0.0008), loss='binary_crossentropy', metrics=['accuracy'])

    return model


if __name__ == '__main__':
    TRAIN_QUEUE = Queue(maxsize=10)
    TRAIN_PROCESS = Process(target=TrainQueueProcess, args=(TRAIN_QUEUE))
    TRAIN_PROCESS.start()

    model = get_model(model_path)

1 Ответ

1 голос
/ 24 марта 2019

Если вы работаете в Windows, переместите все ваши tf и keras импорт в методы.

Как избежать загрузки родительского модуля в разветвленном процессе с многопроцессорной обработкой Pythons

Поскольку в Windows отсутствует os.fork (), все импорта снова импортируются в новом процессе (который в вашем случае включает импорт tf).

https://docs.python.org/2/library/multiprocessing.html#windows

...