使用池的 Tensorflow 错误：无法pickle _thread.RLock 对象

Question

我试图在 Tensorflow CPU 上实现大规模并行微分方程求解器（30k DE），但内存不足（大约 30GB 矩阵）。所以我实现了一个基于批处理的求解器（在短时间内求解并保存数据 -> 设置新的初始值 -> 再次求解）。但问题仍然存在。我了解到在关闭 python 解释器之前，Tensorflow 不会清除内存。因此，基于有关 github 问题的信息，我尝试使用池实现多处理解决方案，但在池化步骤中我不断收到“无法腌制 _thread.RLock 对象”。有人可以帮忙吗！

def dAdt(X,t):

  dX = // vector of differential

  return dX

global state_vector

global state

state_vector =  [0]*n // initial state

def tensor_process():

    with tf.Session() as sess:

        print("Session started...",end="")

        tf.global_variables_initializer().run()

        state = sess.run(tensor_state)

        sess.close()

n_batch = 3

t_batch = np.array_split(t,n_batch)

for n,i in enumerate(t_batch):

    print("Batch",(n+1),"Running...",end="")

    if n>0:

        i = np.append(i[0]-0.01,i)

    print("Session started...",end="")

    init_state = tf.constant(state_vector, dtype=tf.float64)

    tensor_state = tf.contrib.odeint_fixed(dAdt, init_state, i)

    with Pool(1) as p:

        p.apply_async(tensor_process).get()

    state_vector = state[-1,:]

    np.save("state.batch"+str(n+1),state)

    state=None