Tensorflow (Keras) 和多处理导致 GPU 内存不足

Question

我有一个自定义 DataGenerator，它使用 Python 的 Multiprocessing 模块来生成输入到 Tensorflow 模型的训练数据。

问题在于，每当初始化一个新的 DataGenerator 进程时，它似乎都会尝试初始化 Tensorflow（在代码顶部导入）并为自己分配一些 GPU 内存。

我按照this question 限制每个进程对 GPU 内存的访问，我的代码有效，但我只能使用可用 GPU 内存的三分之一。

新流程和 Tensorflow 代码在同一个 Python 文件中启动。是否有适当的方法来利用 Multiprocessing，同时禁止衍生进程导入 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)

Answer 1

如果您在 Windows 上，请将您的所有 tf 和 keras 导入移动到方法中。

How to avoid loading a parent module in a forked process with Pythons multiprocessing

由于 Windows 缺少 os.fork() 所有导入都在新进程中再次导入（在您的情况下包括导入 tf）。

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