0

我使用多处理中的 pool.map 来并行化我的 python 代码。当我用 pool.map 调用我的 tensorflow/keras 模型时，如果我的神经网络大于某个大小，代码就会挂起。我仍然有足够的 RAM 可用，并且在池之外调用模型可以正常工作。

我在 linux 上使用 python 3.7、tensorflow 2.3。

下面提供了一个 mwe，它也在colab 上：

def my_function(i):
    a = MODEL(np.array(i).reshape(1,1))
    print('foo', i)
    return a

THREADS = os.cpu_count()
N = 4
NEURONS = 150000  # works for 100000, hangs for 150000
MODEL = tf.keras.Sequential([tf.keras.layers.Dense(NEURONS, input_shape=(1,))])

my_function(10) # works fine

pool = multiprocessing.Pool(THREADS)
_ = pool.map(my_function, range(N)) # hangs
pool.close()
pool.join()

知道问题是什么吗？如何并行调用大型模型？

编辑： a 的大小不是问题，仅当 tf.keras 在池外调用一次时，代码才会挂起，请参见下面的 mwe 和colab。神经元的临界数量低于原始示例。任何想法？

def my_function(i):
    print('start', i)
    model = tf.keras.Sequential([tf.keras.layers.Dense(NEURONS, input_shape=(1,))])
    print('finish', i)
    return None

THREADS = os.cpu_count()
N = 4
NEURONS = 20000 # works with 10000, not with 20000

# works
pool = multiprocessing.Pool(THREADS)
_ = pool.map(my_function, range(N))
pool.close()
pool.join()

# works
my_function(10)

# doesn't work if many neurons
pool = multiprocessing.Pool(THREADS)
_ = pool.map(my_function, range(N))
pool.close()
pool.join()
4

0 回答 0